Runtime-branchnig for jit-witgen

[chriseth]

If we want to do jit-witgen for dynamic machines containing instructions, we need to implement runtime branching. One simple idea is to take a look at witness columns that are known but not concrete and are constrained to a small number of variables. Branching on all possible values will likely result in progress. The problem is that we will get a number of branches that is exponential in the number of variables. In reality, though, many of these variable-value-combinations exclude each other, especially in the case of instructions.

As an example, take the following constraints that are part of the pil code generated from simple_sum.asm:

X = read_X_A * A + read_X_CNT * CNT + X_read_free * X_free_value;
A' = reg_write_X_A * X + (1 - (reg_write_X_A + instr__reset)) * A;
CNT' = reg_write_X_CNT * X + instr_dec_CNT * (CNT - 1) +
           (1 - (reg_write_X_CNT + instr_dec_CNT + instr__reset)) * CNT;
pc_update = instr_jmpz * (instr_jmpz_pc_update + instr_jmpz_pc_update_1) +
            instr_jmp * instr_jmp_param_l +
            instr__jump_to_operation * _operation_id +
            instr__loop * pc +
            (1 - (instr_jmpz + instr_jmp + instr__jump_to_operation + instr__loop + instr_return)) * (pc + 1);
pc' = (1 - first_step') * pc_update;

By looking a the (1 - (reg_write_X_CNT + ...)) sub-expressions in the CNT' constraint and the (1 - (instr_jmpz + ...) sub-expression in the pc'-constraint, we know that the code was designed such that these summands exclude each other.

This means that for example, if we branch on if instr_jmp = 1, all the other instruction flags can be inferred to be zero and we can directly make progress on pc without the need for further branches and then continue with } else if instr__loop = 1 { etc.

Now how do we derive that the flags are zero?

We need a second stage derivation, where we treat known variables in a slightly different way (and this second-stage derivation could also be useful without branching, by the way):

After having branched on e.g. instr_jmp = 1 (in a certain row!), we set the range constraint of instr_jmp to just the value 1.
Then we go through all identities again, even those that are marked done (even lookups). While evaluating these identities, we mark only those variables as known that are known and concrete, i.e. constrained to a single value. Both the unknown and known-but-not-concrete variables are treated as symbolic unknown variables. From this we derive assignments, which should now always be concrete values. These should be assignments to already known but not concrete variables. We store those new range constraints, but do not generate assignments. Lookups probably need to be treated specially, since we want lookups into the FixedMachine to generate concrete values at compile time.
This operation run on polynomial constraints will propagate values of branched-on variables across rows and columns and run on fixed lookups will generate new concrete values for other columns when they are mutually exclusive (I think the current implementation of the fixed machine will already be suitable).

tl;dr: After we branch and after we make no progress during general witgen derivation, we should add another mode to derive effects where we only consider known-and-concrete variables as known and consider known-and-just-symbolic variables as unknown and thus derive new concrete values for the latter, which should give us new progress in the "regular" derivation.