day7.ml

module Intcode = struct
  type t = {
    program: int array;
    mutable pos: int;
    inputs: int Queue.t;
    mutable outputs: int list;
    mutable last_op: op
  }
  and op = | Halt | Input | Output | Other

  let init ?(inputs=[]) ?(outputs=[]) program =
    let queue = Queue.create () in
    List.iter (fun item -> Queue.push item queue) inputs;
    {
      program = Array.of_list program;
      pos = 0;
      inputs = queue;
      outputs;
      last_op = Other;
    }

  let step (state: t) =
    let get = Array.get state.program in
    let set ~pos = Array.set state.program pos in
    let last_op op = state.last_op <- op in
    let move pos = state.pos <- pos in
    let read () =
      Queue.pop state.inputs
    in
    let write value =
      state.outputs <- value :: state.outputs
    in
    let op = get state.pos |> string_of_int |> CCString.to_list |> List.rev in
    let with_modes modes =
      let modes =
        modes
        |> CCString.of_list
        |> CCString.pad ~side:`Right ~c:'0' 4
        |> CCString.to_list
      in
      let mode param =
        match CCList.drop param modes with
        | [] | '0' :: _ -> `Pos
        | '1' :: _ -> `Value
        | _ -> assert false
      in
      fun param ->
        let value = get (state.pos + param) in
        match mode param with
        | `Value -> value
        | `Pos -> get value
    in
    match op with

    (* halt *)
    | ['9'; '9'] -> last_op Halt;

    (* read *)
    | ['3'] ->
      let input_pos = get (state.pos + 1) in
      let value = read () in
      set ~pos:input_pos value;
      last_op Input;
      move (state.pos + 2)

    (* write *)
    | '4' :: modes ->
      let get_param = with_modes modes in
      write (get_param 1);
      last_op Output;
      move (state.pos + 2)

    (* jump-if-true / jump-if-false *)
    | op :: modes when op = '5' || op = '6' ->
      let get_param = with_modes modes in
      let f = if op = '5' then ( <> ) 0 else ( = ) 0 in
      let next = match get_param 1 |> f with
        | true -> get_param 2
        | false -> state.pos + 3
      in
      last_op Other;
      move next

    (* less-than / equals *)
    | op :: modes when op = '7' || op = '8' ->
      let get_param = with_modes modes in
      let f = if op = '7' then ( < ) else ( = ) in
      let value = match f (get_param 1) (get_param 2) with
        | true -> 1
        | false -> 0
      in
      set ~pos:(get (state.pos + 3)) value;
      last_op Other;
      move (state.pos + 4)

    (* add / mul *)
    | op :: modes when op = '1' || op = '2' ->
      let get_param = with_modes modes in
      let f = if op = '1' then ( + ) else ( * ) in
      let res = f (get_param 1) (get_param 2) in
      let res_pos = get (state.pos + 3) in
      set ~pos:res_pos res;
      last_op Other;
      move (state.pos + 4)

    | op ->
      CCString.of_list op |> print_endline;
      assert false

  let rec run state =
    step state;
    match state.last_op with
    | Halt -> state.outputs
    | _ -> run state

  let rec run_till_output state =
    step state;
    match state.last_op with
    | Halt -> `Halted
    | Output -> `Output
    | _ -> run_till_output state

  let push_input state input =
    Queue.push input state.inputs

  let get_outputs state =
    state.outputs

end

let solve lines =
  let program =
    lines
    |> List.hd
    |> CCString.split_on_char ','
    |> List.map int_of_string
  in
  let max_output = ref min_int in
  let rec exec_on_amp input = function
    | [_; _; _; _; _] ->
      if input > !max_output then max_output := input
    | phases ->
      Iter.(0 -- 4) |> Iter.iter (fun phase ->
          if List.mem phase phases
          then ()
          else
            let state = Intcode.init ~inputs:[phase; input] program in
            let next_input = Intcode.run state |> List.hd in
            exec_on_amp next_input (phase :: phases)
        )
  in
  exec_on_amp 0 [];
  Printf.printf "Part1 = %d\n" !max_output;
  let rec run_chain ?next_input = function
    | hd :: tl ->
      CCOpt.iter (fun input -> Intcode.push_input hd input) next_input;
      let next_input, next_chain =
        match Intcode.run_till_output hd with
        | `Output ->
          Some (hd |> Intcode.get_outputs |> List.hd), (tl @ [hd])
        | `Halted ->
          None, tl
      in
      run_chain ?next_input next_chain
    | [] -> ()
  in
  let open Iter.Infix in
  let rec insert x l = match l with
    | [] -> Iter.return [x]
    | y :: tl ->
      Iter.append
        (insert x tl >|= fun tl' -> y :: tl')
        (Iter.return (x :: l))
  in
  let rec permute = function
    | [] -> Iter.return []
    | x :: tl -> permute tl >>= insert x
  in
  let max_output = ref min_int in
  permute [5;6;7;8;9]
  |> Iter.map (List.map (fun input -> Intcode.init ~inputs:[input] program))
  |> Iter.iter (fun chain ->
      run_chain ~next_input:0 chain;
      let output = chain |> List.rev |> List.hd |> Intcode.get_outputs |> List.hd in
      if output > !max_output then max_output := output
    );
  Printf.printf "Part2 = %d\n" !max_output;