type noun =
  | Atom of Z.t
  | Cell of noun * noun

exception Exit

let ensure_non_negative z =
  if Z.sign z < 0 then raise Exit else z

let atom z =
  Atom (ensure_non_negative z)

let atom_of_int n =
  if n < 0 then raise Exit else Atom (Z.of_int n)

(* Convert ASCII string to atom (bytes in little-endian order) *)
let atom_of_string s =
  if String.length s = 0 then atom Z.zero
  else
    let bytes = Bytes.of_string s in
    atom (Z.of_bits (Bytes.to_string bytes))

let cell h t = Cell (h, t)

let zero = atom_of_int 0
let one = atom_of_int 1

let is_atom = function
  | Atom _ -> true
  | Cell _ -> false

let is_cell = function
  | Cell _ -> true
  | Atom _ -> false

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

let inc = function
  | Atom z -> Atom (Z.succ z)
  | Cell _ -> raise Exit

let head = function
  | Cell (h, _) -> h
  | Atom _ -> raise Exit

let tail = function
  | Cell (_, t) -> t
  | Atom _ -> raise Exit

let rec slot axis noun =
  if Z.equal axis Z.one then noun
  else if Z.equal axis Z.zero then raise Exit
  else
    let bit = Z.testbit axis 0 in
    let parent = Z.shift_right axis 1 in
    let sub = slot parent noun in
    if bit then tail sub else head sub

let rec to_list noun =
  match noun with
  | Atom z when Z.equal z Z.zero -> []
  | Cell (h, t) -> h :: to_list t
  | _ -> raise Exit