include "globals.mzn";

% Types
enum Operation = {
    i_inp,
    i_add,
    i_mul,
    i_div,
    i_mod,
    i_eql
};
enum Pos = { Arg1, Arg2 };
enum Arg = R(Registers) ++ Imm(Immediates) ++ { None };

enum Registers = {x, y, z, w};

set of int: Immediates = -127..127;
set of int: Digits = 0..9;

% Input
array[int] of Operation: inp_ops;
array[int, Pos] of Arg: inp_args;

% Properties derived from input
set of int: Steps = index_set(inp_ops);
set of int: Inputs = 1..count(inp_ops, i_inp);

set of int: ExtendedSteps = min(Steps)..max(Steps) + 1;
set of int: ExtendedInputs = min(Inputs)..max(Inputs) + 1;

% State
array[ExtendedSteps, Registers] of var int: regs;
array[Inputs] of var Digits: inputs;

% Everything is 0 to start with
constraint forall (r in Registers) (
    regs[min(Steps), r] = 0
);

% All registers unaffected by the instruction at a given step
function set of Registers: unaffected(Steps: step) = {
    r | r in Registers where R(r) != inp_args[step, Arg1]
};

% The argument at pos addressed by the instruction at i_step, pointing to the registers at r_step
function var int: arg(Steps: i_step, ExtendedSteps: r_step, Pos: pos) = 
    if inp_args[i_step, pos] in R(Registers) then
        regs[r_step, R^-1(inp_args[i_step, pos])]
    elseif inp_args[i_step, pos] in Imm(Immediates) then
        Imm^-1(inp_args[i_step, pos])
    else
        assert(false, "no value for argument  of instruction at step")
    endif;

% Where to write the result of an instruction
function var int: writearg(Steps: step) = arg(step, step + 1, Arg1);

% Read the argument at pos of instruction at step
function var int: readarg(Steps: step, Pos: pos) = arg(step, step, pos);

predicate interpret(Steps: step, ExtendedInputs: next_input) = (
    % Constrain state based on instruction at step
    if inp_ops[step] = i_inp then
        writearg(step) = inputs[next_input]
    else
        let {
            Operation: o = inp_ops[step],
            var int: a = readarg(step, Arg1),
            var int: b = readarg(step, Arg2),
            var int: x = writearg(step),
        } in 
        if o = i_add then
            x = a + b
        elseif o = i_mul then
            x = a * b
        elseif o = i_div then
            b != 0 /\
            x * b = a - (a mod b)
        elseif o = i_mod then
            a >= 0 /\
            b > 0 /\
            x = a mod b
        elseif o = i_eql then
            x = (a = b)
        else
            assert(false, "unimplemented instruction at ")
        endif
    endif /\
    % Leave the rest unchanged
    forall (r in unaffected(step)) (
        regs[step + 1, r] = regs[step, r]
    ) /\
    % Perform induction
    if step + 1 in Steps then
        interpret(step + 1, next_input + (inp_ops[step] = i_inp))
    else 
        true
    endif
);

constraint interpret(min(Steps), min(Inputs));

% Finish with z = 0
constraint regs[max(ExtendedSteps), z] = 0;

% No 0s in serial number
constraint forall (i in Inputs) (
    inputs[i] != 0
);

% Find the largest first
solve :: int_search(inputs, input_order, indomain_max) satisfy;

output [
  "inputs:          ", show(inputs), "\n",
  "final registers: ", show(regs[max(ExtendedSteps), ..]), "\n",
];