(** Noun type and basic operations *)

(** A noun is either an atom (arbitrary-precision integer) or a cell (pair of nouns) *)
type noun =
  | Atom of Z.t  (** Arbitrary-precision integer using Zarith *)
  | Cell of noun * noun  (** Pair of nouns *)

(** Exception raised on nock evaluation errors *)
exception Exit

(** Create an atom from an int *)
let atom n = Atom (Z.of_int n)

(** Create a cell *)
let cell a b = Cell (a, b)

(** Test if a noun is a cell *)
let is_cell = function
  | Cell _ -> true
  | Atom _ -> false

(** Test if a noun is an atom *)
let is_atom = function
  | Atom _ -> true
  | Cell _ -> false

(** Get head of a cell *)
let head = function
  | Cell (h, _) -> h
  | Atom _ -> raise Exit

(** Get tail of a cell *)
let tail = function
  | Cell (_, t) -> t
  | Atom _ -> raise Exit

(** Fragment/axis lookup: slot(n, noun)
    This implements the tree-addressing scheme:
    - 1 is the root
    - 2 is head, 3 is tail
    - For n > 1: if even, go left; if odd, go right
*)
let slot_debug = ref false

let rec slot n noun =
  let debug = !slot_debug && (Z.equal n (Z.of_int 2) || Z.equal n (Z.of_int 3)) in
  if debug then
    Printf.eprintf "[SLOT-DEBUG] slot(%s, noun)\n%!" (Z.to_string n);

  if Z.equal n Z.one then begin
    if debug then Printf.eprintf "[SLOT-DEBUG] slot(%s) = identity\n%!" (Z.to_string n);
    noun
  end
  else if Z.equal n Z.zero then
    raise Exit
  else begin
    let bit = Z.testbit n 0 in  (* Check if odd *)
    let parent = Z.shift_right n 1 in
    if debug then
      Printf.eprintf "[SLOT-DEBUG] slot(%s): bit=%b parent=%s\n%!"
        (Z.to_string n) bit (Z.to_string parent);
    let sub = slot parent noun in
    let result = if bit then tail sub else head sub in
    if debug then begin
      Printf.eprintf "[SLOT-DEBUG] slot(%s): taking %s, result is %s\n%!"
        (Z.to_string n) (if bit then "tail" else "head")
        (if is_cell result then
          let h = head result in
          if is_atom h then
            (match h with Atom z -> "cell[" ^ Z.to_string z ^ " ...]" | _ -> "cell[? ...]")
          else "cell[cell ...]"
        else "atom")
    end;
    result
  end

(** Equality test for nouns *)
let rec equal a b =
  match a, b with
  | Atom x, Atom y -> Z.equal x y
  | Cell (ah, at), Cell (bh, bt) -> equal ah bh && equal at bt
  | _, _ -> false

(** Increment an atom *)
let inc = function
  | Atom n -> Atom (Z.succ n)
  | Cell _ -> raise Exit

(** Pretty-print a noun *)
let rec pp_noun fmt = function
  | Atom n -> Format.fprintf fmt "%s" (Z.to_string n)
  | Cell (a, b) -> Format.fprintf fmt "[%a %a]" pp_noun a pp_noun b