brainfuck.asm

machine Brainfuck {
	reg pc[@pc];
	reg X[<=];
	reg Y[<=];
	reg Z[<=];

	// The pc of the given brainfuck program
	reg b_pc;
	// The current operator
	reg op;
	// Data pointer
	reg dp;
	// program's stdin counter
	reg in_ptr;

	// General purpose registers
	reg x1;
	reg t1;
	reg t2;

    instr jump l: label -> Y { pc' = l, Y = pc + 1}
    instr jump_dyn X -> Y { pc' = X, Y = pc + 1}
    instr branch_if_zero X, l: label { pc' = XIsZero * l + (1 - XIsZero) * (pc + 1) }
	instr fail { 1 = 0 }

	function main {
		x1 <== jump(__runtime_start);

		exit:
			return;

		run_op:
			// '>' is 62
			branch_if_zero op - 62, routine_move_right;
			// '<' is 60
			branch_if_zero op - 60, routine_move_left;
			// '+' is 43
			branch_if_zero op - 43, routine_inc;
			// '-' is 45
			branch_if_zero op - 45, routine_dec;
			// ',' is 44
			branch_if_zero op - 44, routine_read;
			// '.' is 46
			branch_if_zero op - 46, routine_write;
			// unknown op
			fail;
			
		routine_move_right:
			dp <=X= dp + 4;
			t2 <== jump(end_run_op);

		routine_move_left:
			dp <=X= dp - 4;
			t2 <== jump(end_run_op);

		routine_inc:
			t1, t2 <== mload(dp);
			mstore dp, t1 + 1;
			t2 <== jump(end_run_op);

		routine_dec:
			t1, t2 <== mload(dp);
			mstore dp, t1 - 1;
			t2 <== jump(end_run_op);

		routine_read:
			t1 <=X= ${ std::prover::Query::Input(std::convert::int(std::prover::eval(in_ptr))) };
			in_ptr <=X= in_ptr + 1;
			mstore dp, t1;
			t2 <== jump(end_run_op);

		routine_write:
			t1, t2 <== mload(dp);
			t1 <=X= ${ std::prover::Query::PrintChar(std::convert::int(std::prover::eval(dp))) };
			t2 <== jump(end_run_op);

		end_run_op:
			t2 <== jump_dyn(x1);

		__runtime_start:
			// The input is formed in this way:
			// <program_length> <program> <input>
			t1 <=X= ${ std::prover::Query::Input(0) };
			in_ptr <=X= t1 + 1;
			b_pc <=X= 1;
			dp <=X= 40;

		interpreter_loop:
			// We read the program from fd 0
			// TODO this can be optimized by reading the whole program at once to memory
			// TODO we should also hash the program and expose as public
			op <=X= ${ std::prover::Query::Input(std::convert::int(std::prover::eval(b_pc))) };

			branch_if_zero op, exit;

			b_pc <=X= b_pc + 1;
			x1 <== jump(run_op);

			t2 <== jump(interpreter_loop);
	}

    // ============== memory instructions ==============

    let up_to_three: col = |i| i % 4;
    let six_bits: col = |i| i % 2**6;
 
    instr mload Y -> X, Z {
        // Z * (Z - 1) * (Z - 2) * (Z - 3) = 0,
        { Z } in { up_to_three },
        Y = wrap_bit * 2**32 + X_b4 * 0x1000000 + X_b3 * 0x10000 + X_b2 * 0x100 + X_b1 * 4 + Z,
        { X_b1 } in { six_bits },
        {
            X_b4 * 0x1000000 + X_b3 * 0x10000 + X_b2 * 0x100 + X_b1 * 4,
            STEP,
            X
        } is m_is_read { m_addr, m_step, m_value }
        // If we could access the shift machine here, we
        // could even do the following to complete the mload:
        // { W, X, Z} in { shr.value, shr.amount, shr.amount}
    }

    /// Stores Z at address Y % 2**32. Y can be between 0 and 2**33.
    /// Y should be a multiple of 4, but this instruction does not enforce it.
    instr mstore Y, Z {
        { X_b1 + X_b2 * 0x100 + X_b3 * 0x10000 + X_b4 * 0x1000000, STEP, Z } is m_is_write { m_addr, m_step, m_value },
        // Wrap the addr value
        Y = (X_b1 + X_b2 * 0x100 + X_b3 * 0x10000 + X_b4 * 0x1000000) + wrap_bit * 2**32
    }

	// =============== read-write memory =======================
    // Read-write memory. Columns are sorted by m_addr and
    // then by m_step. m_change is 1 if and only if m_addr changes
    // in the next row.
    col witness m_addr;
    col witness m_step;
    col witness m_change;
    col witness m_value;

    // Memory operation flags
    col witness m_is_write;
    col witness m_is_read;

    // All operation flags are boolean and either all 0 or exactly 1 is set.
    std::utils::force_bool(m_is_write);
    std::utils::force_bool(m_is_read);
    m_is_read * m_is_write = 0;

    // If the next line is a not a write and we have an address change,
    // then the value is zero.
    (1 - m_is_write') * m_change * m_value' = 0;

    // m_change has to be 1 in the last row, so that a first read on row zero is constrained to return 0
    (1 - m_change) * LAST = 0;

    // If the next line is a read and we stay at the same address, then the
    // value cannot change.
    (1 - m_is_write') * (1 - m_change) * (m_value' - m_value) = 0;

    col witness m_diff_lower;
    col witness m_diff_upper;

    col fixed FIRST = [1] + [0]*;
    let LAST = FIRST';
    col fixed STEP(i) { i };
    col fixed BIT16(i) { i & 0xffff };

    {m_diff_lower} in {BIT16};
    {m_diff_upper} in {BIT16};

    std::utils::force_bool(m_change);

    // if m_change is zero, m_addr has to stay the same.
    (m_addr' - m_addr) * (1 - m_change) = 0;

    // Except for the last row, if m_change is 1, then m_addr has to increase,
    // if it is zero, m_step has to increase.
    // `m_diff_upper * 2**16 + m_diff_lower` has to be equal to the difference **minus one**.
    // Since we know that both m_addr and m_step can only be 32-Bit, this enforces that
    // the values are strictly increasing.
    col diff = (m_change * (m_addr' - m_addr) + (1 - m_change) * (m_step' - m_step));
    (1 - LAST) * (diff - 1 - m_diff_upper * 2**16 - m_diff_lower) = 0;

	col fixed bytes(i) { i & 0xff };
    col witness X_b1;
    col witness X_b2;
    col witness X_b3;
    col witness X_b4;
    { X_b1 } in { bytes };
    { X_b2 } in { bytes };
    { X_b3 } in { bytes };
    { X_b4 } in { bytes };
    col witness wrap_bit;
    wrap_bit * (1 - wrap_bit) = 0;

    // Input is a 32 bit unsigned number. We check bit 7 and set all higher bits to that value.
    instr sign_extend_byte Y -> X {
        // wrap_bit is used as sign_bit here.
        Y = Y_7bit + wrap_bit * 0x80 + X_b2 * 0x100 + X_b3 * 0x10000 + X_b4 * 0x1000000,
        X = Y_7bit + wrap_bit * 0xffffff80
    }
    col fixed seven_bit(i) { i & 0x7f };
    col witness Y_7bit;
    { Y_7bit } in { seven_bit };

    // ============== iszero check for X =======================
    let XIsZero = std::utils::is_zero(X);
}