✅Collatz Puzzle

Idea

The parameter addr should be a contract that we deploy, so this is our solution contract. The function callMe() takes addr as input and does some math. In the following line:

q = ICollatz(addr).collatzIteration{gas: 100}(q);

We learn that it calls the collatzIteration() function in our solution contract, so our task is to implement this function with respect to the following standard:

  function collatzIteration(uint256 n) public pure override returns (uint256) {
    if (n % 2 == 0) {
      return n / 2;
    } else {
      return 3 * n + 1;
    }
  }

There are two difficulties we encounter:

1. Our solution contract must have size ≤ 32 bytes (size := extcodesize(addr)).

2. The gas limit is 100 (ICollatz(addr).collatzIteration{gas: 100}(q))

Therefore, there is no way we can do this with Solidity, even with assembly. In such scenario, Huff is a great choice since it is low-level enough for optimization.

In Huff, usually we want to implement things like function dispatcher. For this challenge, we will skip that part since the code size is heavily constrained. Skipping function dispatching works because we can implement the logic of collatzIteration() inside MAIN macro and simply throw away the first 4 bytes in calldata. Since the MAIN macro is the entry point, the logic inside it will be executed anyway. That is, calling collatzIteration() is the same as calling the MAIN macro.

Just keep in mind that skipping function dispatching is bad practice in production.

First we write a Foundry test for callMe():

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.15;

import "foundry-huff/HuffDeployer.sol";
import "forge-std/Test.sol";
import "forge-std/console.sol";
import {CollatzPuzzle} from "src/CollatzPuzzle.sol";

interface ICollatz {
  function collatzIteration(uint256 n) external pure returns (uint256);
}

contract CollatzPuzzleSolutionTest is Test {
    CollatzPuzzle public puzzle;
    ICollatz public solution;

    function setUp() public {
        // This is the caller
        puzzle = new CollatzPuzzle();
        // This is the callee
        solution = ICollatz(HuffDeployer.deploy("Solution"));
    }

    /// @dev Ensure that you can set and get the value.
    function testCallMe() public {
        address addr = address(solution);

        // Log Huff code size
        uint256 size;
        assembly {
            size := extcodesize(addr)
        }
        console.log("Code size: ", size);

        puzzle.callMe(addr);
    }
}

This piece of code grabs the Huff code size and display it in the console:

        // Log Huff code size
        uint256 size;
        assembly {
            size := extcodesize(addr)
        }
        console.log("Code size: ", size);

In order to read the content of console.log(), we should use -vv verbosity when running forge test.

Now let’s implement the logic of collatzIteration(). Here is my first try:

/*
  function collatzIteration(uint256 n) public pure override returns (uint256) {
    if (n % 2 == 0) {
      return n / 2;
    } else {
      return 3 * n + 1;
    }
  }
*/

// Function dispatching skipped for compact code size
#define macro MAIN() = takes (0) returns (0) 
    0x04 calldataload // [n]

    // Push an extra n onto the stack for future computation
    0x02 dup2 // [n, 2, n]

    // Test if n % 2 == 0
    mod // [n % 2, n]
    iszero // [n % 2 == 0, n]

    // Jump to the "even" case if n % 2 == 0
    even jumpi

    // Otherwise, n is odd
    0x03 mul // [3 * n]
    0x01 add // [3 * n + 1]
    ret jump
    
    // Handle the case when n is even
    even:
        0x02 swap1 div // [n / 2]
        ret jump
    // Return logic
    ret:
        0x00 mstore // Store the top element from the stack in memory offset 0x00
        0x20 0x00 return // Return 32 bytes from offset 0x00
}

Test result returns “bad code size!” and shows that the current code size is 40:

Oversize

Let’s see how to optimize the Huff code.

Take a look at this part of the code:

    // Handle the case when n is even
    even:
        0x02 swap1 div // [n / 2]
        ret jump
    // Return logic
    ret:
        0x00 mstore // Store the top element from the stack in memory offset 0x00
        0x20 0x00 return // Return 32 bytes from offset 0x00

The ret jump here is redundant since the control flow will move on to the ret label anyway. Deleting ret jump makes code size 36. We are 4 bytes away from our goal.

The next optimization is a bit tricky: we can actually get rid of the ret jump at the end of the odd case as well.

Here is the logic of the odd case:

    // Otherwise, n is odd
    0x03 mul // [3 * n]
    0x01 add // [3 * n + 1]
    ret jump
    
    // Handle the case when n is even
    even:
        0x02 swap1 div // [n / 2]
    // Return logic
    ret:
        0x00 mstore // Store the top element from the stack in memory offset 0x00
        0x20 0x00 return // Return 32 bytes from offset 0x00

Here we can compute 2 * (3 * n + 1) in the odd case and let the control flow goes into the even case and get 2 * (3 * n + 1) / 2, and the computation is still correct. Modified version:

    // Otherwise, n is odd
    0x03 mul // [3 * n]
    0x01 add // [3 * n + 1]
    0x02 mul // [2 * (3 * n + 1)]
    
    // Handle the case when n is even
    even:
        0x02 swap1 div // [n / 2]
    // Return logic
    ret:
        0x00 mstore // Store the top element from the stack in memory offset 0x00
        0x20 0x00 return // Return 32 bytes from offset 0x00

This gives code size 35. The math expression in the odd case can be simplified to 6 * n + 2, so we come up with a more optimized version:

    // Otherwise, n is odd
    0x06 mul // [6 * n]
    0x02 add // [6 * n + 2]
    
    // Handle the case when n is even
    even:
        0x02 swap1 div // [n / 2]
    // Return logic
    ret:
        0x00 mstore // Store the top element from the stack in memory offset 0x00
        0x20 0x00 return // Return 32 bytes from offset 0x00

This version gives code size exactly 32 and the test is passed:

Test passed

PoC

GitHub - ret2basic/Collatz-Puzzle: QuillCTF "Collatz Puzzle" solution.GitHub

Collatz Puzzle