InjectCode.sol

// SPDX-License-Identifier: GPL-3.0

pragma solidity >=0.7.0 <0.9.0;

/**
 * Workaround example on how to inject and execute arbitrary bytecode in solidity contract
 * Currently only YUL supports verbatim: https://github.com/ethereum/solidity/issues/12067
 * But you cannot import Solidity code in YUL, or YUL code in solidity, so this workaround is necessary.
 * It works as long the byte sequence `0x7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F00` appear in the runtime code.
 *
 * for testing you must deploy the `Example` contract, not the `ExampleImpl`
 */
library CodeInjection {
    struct Builder {
        bytes bytecode;
        uint256 mask;
        bytes32[256] code;
    }

    function fromBytecode(bytes memory bytecode) internal pure returns (Builder memory r) {
        bytes32[256] memory code;
        r = Builder({
            bytecode: bytecode,
            mask: 0,
            code: code
        });
    }

    function inject(Builder memory bytecode, uint8 tag, bytes32 code) internal pure {
        bytecode.mask |= 1 << uint256(tag);
        bytecode.code[tag] = code;
    }

    /**
     * @dev Replaces all occurences of `0x7F7F7F...tag` by the respective injected bytecode
     **/
    function build(Builder memory bytecode) internal pure returns (bytes memory finalBytecode) {
        bytes32[256] memory codes = bytecode.code;
        assembly ("memory-safe") {
            for {
                let pos := mload(bytecode)
                let end := mload(pos)
                pos := add(pos, 0x20)
                end := add(end, pos)
                let mask := mload(add(bytecode, 0x20))
            } mul(mask, lt(pos, end)) {} {
                // Efficient Algorithm to find 32 consecutive repeated bytes in a byte sequence
                // It look in chunks of 32 bytes, and works even if the constant is not aligned.
                for {
                    let chunk := 1
                } gt(chunk, 0) { pos := add(pos, chunk) } {
                    // Transform all `0x7F` bytes into `0xFF`
                    // 0x80 ^ 0x7F == 0xFF
                    // Also transform all other bytes in something different than `0xFF`
                    chunk := xor(mload(pos), 0x8080808080808080808080808080808080808080808080808080808080808080)

                    // Find the right most unset bit, which is equivalent to find the
                    // right most byte different than `0x7F`.
                    // ex: (0x12345678FFFFFF + 1) & (~0x12345678FFFFFF) == 0x00000001000000
                    chunk := and(add(chunk, 1), not(chunk))

                    // Round down to the closest multiple of 256
                    // Ex: 2 ** 18 become 2 ** 16
                    chunk := div(chunk, mod(chunk, 0xff))

                    // Find the number of leading bytes different than `0x7E`.
                    // Rationale:
                    // Multiplying a number by a power of 2 is the same as shifting the bits to the left
                    // 1337 * (2 ** 16) == 1337 << 16
                    // Once the chunk is a multiple of 256 it always shift entire bytes, we use this to
                    // select a specific byte in a byte sequence.
                    chunk := byte(0, mul(0x201f1e1d1c1b1a191817161514131211100f0e0d0c0b0a090807060504030201, chunk))

                    // Stop the loop if we go out of bounds
                    // obs: can remove this check if you are 100% sure the constant exists
                    chunk := mul(chunk, lt(pos, end))
                }

                let tag := and(mload(add(pos, 1)), 0xff)
                let offset := add(codes, shl(5, tag))
                let code := mload(offset)
                if or(iszero(code), iszero(lt(pos, end))) {
                    revert(0, 0)
                }
                mask := and(mask, not(shl(tag, 1)))
                mstore(add(pos, 1), 0x5B)
                mstore(pos, code)
                mstore(offset, 0)
                pos := add(pos, 32)
            }

            finalBytecode := mload(bytecode)
        }
    }
}

/**
 * @dev Implementation, should not deploy this one
 */
contract ExampleImpl {
    function add(uint256, uint256) external pure returns (uint256) {
        // tag: 0x01
        return 0x7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F01;
    }

    function sub(uint256, uint256) external pure returns (uint256) {
        // tag: 0x02
        return 0x7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F02;
    }

    function mul(uint256, uint256) external pure returns (uint256) {
        // tag: 0x03
        return 0x7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F03;
    }

    function div(uint256, uint256) external pure returns (uint256) {
        // tag: 0x04
        return 0x7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F7F04;
    }
}

/**
 * @dev Injected implementation, must deploy this one
 */
contract Example is ExampleImpl  {
    using CodeInjection for CodeInjection.Builder;

    // OBS: This codes replaces the PUSH31 OPCODE, so it MUST have exact 32 bytes in size (31 + opcode
    // with no inputs and push ONE value onto the stack.

    // PUSH26 0x04 CALLDATALOAD PUSH1 0x24 CALLDATALOAD ADD
    bytes32 private constant ADD_BYTECODE = 0x7900000000000000000000000000000000000000000000000000043560243501;
    // PUSH26 0x24 CALLDATALOAD PUSH1 0x04 CALLDATALOAD SUB
    bytes32 private constant SUB_BYTECODE = 0x7900000000000000000000000000000000000000000000000000243560043503;
    // PUSH26 0x24 CALLDATALOAD PUSH1 0x04 CALLDATALOAD MUL
    bytes32 private constant MUL_BYTECODE = 0x7900000000000000000000000000000000000000000000000000243560043502;
    // PUSH26 0x24 CALLDATALOAD PUSH1 0x04 CALLDATALOAD DIV
    bytes32 private constant DIV_BYTECODE = 0x7900000000000000000000000000000000000000000000000000243560043504;

    constructor() payable {
        // In solidity the child's constructor are executed before the parent's constructor,
        // so once this contract extends `ExampleImpl`, it's constructor is executed first.

        // Copy `ExampleImpl` runtime code into memory.
        CodeInjection.Builder memory builder = CodeInjection.fromBytecode(type(ExampleImpl).runtimeCode);

        // Inject code
        builder.inject(0x01, ADD_BYTECODE);
        builder.inject(0x02, SUB_BYTECODE);
        builder.inject(0x03, MUL_BYTECODE);
        builder.inject(0x04, DIV_BYTECODE);

        // Build injected bytecode
        bytes memory bytecode = builder.build();

        // Return the new bytecode
        assembly ("memory-safe") {
            // Return the modified bytecode
            return (add(bytecode, 32), mload(bytecode))   
        }
    }
}