module MenhirFolder(M: MenhirLib.IncrementalEngine.INCREMENTAL_ENGINE)(F : sig
    type 'a state
    type 'a frame
    val reduce : 'a state -> M.production -> pop:'a frame -> goto:'a M.env -> 'a frame * 'a state
    val shift : 'a state -> 'a M.env -> 'a frame * 'a state
  end): sig
  open F
  type 'a t
  val make : 'a M.checkpoint -> 'a frame -> 'a state -> 'a t
  val update : 'a t -> 'a M.checkpoint -> 'a t option
  val observe : 'a t -> 'a state
end =
struct
  open F

  type 'a t =
    | Normal of {
        state: 'a state;
        stack: ('a M.env * 'a frame) list }
    | Reduction of {
        state: 'a state;
        stack: ('a M.env * 'a frame) list;
        prod : M.production;
      }

  let make cp frame state =
    match cp with
    | M.InputNeeded env -> Normal { state; stack = [env, frame] }
    | _ -> invalid_arg
             "MenhirFolder: initial checkpoint should be in InputNeeded state"

  let rec pop_until stack env =
    match stack with
    | [] -> assert false
    | (env', f') :: _ when M.equal env env' -> f', stack
    | _ :: tl -> pop_until tl env

  let reduce_to state stack prod env =
    match M.pop env with
    | Some env' ->
      let frame, stack = pop_until stack env' in
      let frame', state' = reduce state prod ~pop:frame ~goto:env in
      state', (env, frame') :: stack
    | None -> assert false

  let get_stack t env =
    match t with
    | Normal { state; stack } -> state, stack
    | Reduction { state; stack; prod } -> reduce_to state stack prod env

  let update t = function
    | M.HandlingError _ | M.Rejected | M.Accepted _ -> None
    | M.InputNeeded env ->
      let state, stack = get_stack t env in
      Some (Normal {state; stack})
    | M.AboutToReduce (env, prod) ->
      let state, stack = get_stack t env in
      Some (Reduction {state; stack; prod})
    | M.Shifting (env, env', _) ->
      let state0, stack = get_stack t env in
      let frame', state = shift state0 env' in
      let stack = (env', frame') :: stack in
      Some (Normal {state; stack})

  let observe (Normal {state; _} | Reduction {state; _}) = state
end


type lr1_state = int
type production = int

type frame = {
  origin: lr1_state;
  reductions: int list;
}

type cell =
  | Shift of lr1_state
  | Reduce of production * cell

module Recorder(M: MenhirLib.IncrementalEngine.INCREMENTAL_ENGINE) = struct

  module Tracker : sig
    type 'a state = {
      transitions: (cell * lr1_state, int ref) Hashtbl.t;
      cell: cell;
    }

    type 'a frame = cell

    val reduce : 'a state -> M.production -> pop:'a frame -> goto:'a M.env -> 'a frame * 'a state
    val shift : 'a state -> 'a M.env -> 'a frame * 'a state
  end = struct
    type 'a state = {
      transitions: (cell * lr1_state, int ref) Hashtbl.t;
      cell: cell;
    }

    type 'a frame = cell

    let incr transitions cell target =
      let key = (cell, target) in
      match Hashtbl.find transitions key with
      | r -> incr r
      | exception Not_found ->
        Hashtbl.add transitions key (ref 1)

    let incr_state state target =
      incr state.transitions state.cell target

    let shift state env =
      let target = M.current_state_number env in
      incr_state state target;
      let cell = Shift target in
      cell, {state with cell}

    let reduce state prod ~pop ~goto =
      let target = M.current_state_number goto in
      incr state.transitions pop target;
      let cell = Reduce (M.production_index prod, state.cell) in
      cell, {state with cell}
  end
end