(** Fully iterative Nock interpreter using explicit stack

    No recursion at all - uses a work stack and result stack.
    This will never overflow regardless of computation depth.
*)

open Noun

type frame =
  | EvalFrame of noun * noun                  (* Need to evaluate nock_on(bus, fol) *)
  | DistLeft of noun * noun                   (* After computing left of [a b], compute right *)
  | DistBoth of noun                          (* After both sides, construct cell *)
  | Op2Formula of noun * noun                 (* After formula, compute subject *)
  | Op2Subject of noun                        (* After subject, tail-call *)
  | Op3Apply                                  (* Apply cell test *)
  | Op4Apply                                  (* Apply increment *)
  | Op5Apply                                  (* Apply equality *)
  | Op6Test of noun * noun * noun             (* After test, select branch (bus, c, d) *)
  | Op7Subject of noun                        (* After subject, tail-call with formula *)
  | Op8Pin of noun * noun                     (* After pin, extend and tail-call (old_bus, c) *)
  | Op9Core of noun * noun                    (* After core, extract slot (bus, b_gal/axis) *)

let nock_on init_bus init_fol =
  let work_stack = ref [EvalFrame (init_bus, init_fol)] in
  let result_stack = ref [] in

  let rec loop () =
    match !work_stack with
    | [] ->
        (match !result_stack with
        | [result] -> result
        | _ -> raise Exit)  (* Should have exactly one result *)

    | EvalFrame (bus, fol) :: work_rest ->
        work_stack := work_rest;

        (match fol with
        | Cell { h = hib; t = gal; _ } when is_cell hib ->
            (* Distribution: push frames in reverse order *)
            work_stack := EvalFrame (bus, hib) :: DistLeft (bus, gal) :: !work_stack;
            loop ()

        | Cell { h = Atom { z = op; _ }; t = gal; _ } ->
            if Z.compare op (Z.of_int max_int) > 0 then raise Exit;

            (match Z.to_int op with
            | 0 ->
                if not (is_atom gal) then raise Exit;
                let axis = (match gal with Atom { z = n; _ } -> n | _ -> raise Exit) in
                let res = slot axis bus in
                result_stack := res :: !result_stack;
                loop ()

            | 1 ->
                result_stack := gal :: !result_stack;
                loop ()

            | 2 ->
                if not (is_cell gal) then raise Exit;
                let c_gal = tail gal in
                let b_gal = head gal in
                work_stack := EvalFrame (bus, c_gal) :: Op2Formula (bus, b_gal) :: !work_stack;
                loop ()

            | 3 ->
                work_stack := EvalFrame (bus, gal) :: Op3Apply :: !work_stack;
                loop ()

            | 4 ->
                work_stack := EvalFrame (bus, gal) :: Op4Apply :: !work_stack;
                loop ()

            | 5 ->
                work_stack := EvalFrame (bus, gal) :: Op5Apply :: !work_stack;
                loop ()

            | 6 ->
                if not (is_cell gal) then raise Exit;
                let b_gal = head gal in
                let cd_gal = tail gal in
                if not (is_cell cd_gal) then raise Exit;
                let c_gal = head cd_gal in
                let d_gal = tail cd_gal in
                work_stack := EvalFrame (bus, b_gal) :: Op6Test (bus, c_gal, d_gal) :: !work_stack;
                loop ()

            | 7 ->
                if not (is_cell gal) then raise Exit;
                let b_gal = head gal in
                let c_gal = tail gal in
                work_stack := EvalFrame (bus, b_gal) :: Op7Subject c_gal :: !work_stack;
                loop ()

            | 8 ->
                if not (is_cell gal) then raise Exit;
                let b_gal = head gal in
                let c_gal = tail gal in
                work_stack := EvalFrame (bus, b_gal) :: Op8Pin (bus, c_gal) :: !work_stack;
                loop ()

            | 9 ->
                if not (is_cell gal) then raise Exit;
                let b_gal = head gal in
                let c_gal = tail gal in
                if not (is_atom b_gal) then raise Exit;
                work_stack := EvalFrame (bus, c_gal) :: Op9Core (bus, b_gal) :: !work_stack;
                loop ()

            | 10 ->
                (* Hint - just evaluate tail, ignore head *)
                if not (is_cell gal) then raise Exit;
                let q_gal = tail gal in
                work_stack := EvalFrame (bus, q_gal) :: !work_stack;
                loop ()

            | 11 ->
                if not (is_cell gal) then raise Exit;
                let q_gal = tail gal in
                work_stack := EvalFrame (bus, q_gal) :: !work_stack;
                loop ()

            | _ -> raise Exit)

        | _ -> raise Exit)

    | DistLeft (bus, gal) :: work_rest ->
        (match !result_stack with
        | left :: result_rest ->
            result_stack := result_rest;
            work_stack := EvalFrame (bus, gal) :: DistBoth left :: work_rest;
            loop ()
        | _ -> raise Exit)

    | DistBoth left :: work_rest ->
        (match !result_stack with
        | right :: result_rest ->
            result_stack := cell left right :: result_rest;
            work_stack := work_rest;
            loop ()
        | _ -> raise Exit)

    | Op2Formula (bus, b_gal) :: work_rest ->
        (match !result_stack with
        | nex :: result_rest ->
            result_stack := result_rest;
            work_stack := EvalFrame (bus, b_gal) :: Op2Subject nex :: work_rest;
            loop ()
        | _ -> raise Exit)

    | Op2Subject nex :: work_rest ->
        (match !result_stack with
        | seb :: result_rest ->
            result_stack := result_rest;
            (* Tail call: push new eval frame *)
            work_stack := EvalFrame (seb, nex) :: work_rest;
            loop ()
        | _ -> raise Exit)

    | Op3Apply :: work_rest ->
        (match !result_stack with
        | gof :: result_rest ->
            let res = if is_cell gof then atom 0 else atom 1 in
            result_stack := res :: result_rest;
            work_stack := work_rest;
            loop ()
        | _ -> raise Exit)

    | Op4Apply :: work_rest ->
        (match !result_stack with
        | gof :: result_rest ->
            result_stack := inc gof :: result_rest;
            work_stack := work_rest;
            loop ()
        | _ -> raise Exit)

    | Op5Apply :: work_rest ->
        (match !result_stack with
        | wim :: result_rest ->
            if not (is_cell wim) then raise Exit;
            let a = head wim in
            let b = tail wim in
            let res = if equal a b then atom 0 else atom 1 in
            result_stack := res :: result_rest;
            work_stack := work_rest;
            loop ()
        | _ -> raise Exit)

    | Op6Test (bus, c_gal, d_gal) :: work_rest ->
        (match !result_stack with
        | tys :: result_rest ->
            result_stack := result_rest;
            let nex = match tys with
              | Atom { z = n; _ } when Z.equal n Z.zero -> c_gal
              | Atom { z = n; _ } when Z.equal n Z.one -> d_gal
              | _ -> raise Exit
            in
            work_stack := EvalFrame (bus, nex) :: work_rest;
            loop ()
        | _ -> raise Exit)

    | Op7Subject c_gal :: work_rest ->
        (match !result_stack with
        | bod :: result_rest ->
            result_stack := result_rest;
            work_stack := EvalFrame (bod, c_gal) :: work_rest;
            loop ()
        | _ -> raise Exit)

    | Op8Pin (old_bus, c_gal) :: work_rest ->
        (match !result_stack with
        | heb :: result_rest ->
            result_stack := result_rest;
            let new_bus = cell heb old_bus in
            work_stack := EvalFrame (new_bus, c_gal) :: work_rest;
            loop ()
        | _ -> raise Exit)

    | Op9Core (_bus, b_gal) :: work_rest ->
        (match !result_stack with
        | cor :: result_rest ->
            result_stack := result_rest;
            if not (is_atom b_gal) then raise Exit;
            let axis = (match b_gal with Atom { z = n; _ } -> n | _ -> raise Exit) in
            let arm = slot axis cor in
            work_stack := EvalFrame (cor, arm) :: work_rest;
            loop ()
        | _ -> raise Exit)

  in

  loop ()