(******************************************************************************)
(*                                                                            *)
(*                                    Menhir                                  *)
(*                                                                            *)
(*   Copyright Inria. All rights reserved. This file is distributed under     *)
(*   the terms of the GNU General Public License version 2, as described in   *)
(*   the file LICENSE.                                                        *)
(*                                                                            *)
(******************************************************************************)

open MiddleAPI

(* This module does not need much functionality: just some basic operations
   on the type [Symbol.t] and on the type [_ Item.Map.t]. For simplicity,
   I require much more, by requesting [G : GRAMMAR]. *)

module Make (G : GRAMMAR) = struct
open G

(* -------------------------------------------------------------------------- *)

(* A level in a forest consists in:
   - a prefix (a list of symbols),
   - a tree that is in focus,
   - a suffix (a list of symbols),
   - a comment. *)

type 'focus level = {
    prefix: Symbol.t list;
    focus: 'focus;
    suffix: Symbol.t list;
    comment: string
  }

(* A tree either is empty or consists of a root symbol and a child forest. *)

type tree =
  | TEmpty
  | TRooted of Symbol.t * forest

(* A forest is a level with a tree in focus. *)

and forest =
    tree level

(* We make use of two kinds of contexts with a forest-shaped hole.
   We have tree contexts and forest contexts.

   Tree contexts do not have a case for holes, because we work with
   forest-shaped holes only. Forest contexts have one. *)

(* A tree context consists of a root symbol and a child forest context. *)

type ctree =
  | CRooted of Symbol.t * cforest

(* A forest context is either a hole
   or a forest with a tree context in focus. *)

and cforest =
  | CHole
  | CCons of ctree level

(* The public type [t] represents a forest.
   The public type [context] represents a forest context. *)

type t =
  forest

type context =
  cforest

(* -------------------------------------------------------------------------- *)

(* Construction. *)

let empty =
  {
    prefix = [];
    focus = TEmpty;
    suffix = [];
    comment = ""
  }

let tail i rhs =
  let n = Array.length rhs in
  (* The case [i = n] is forbidden because this function is useful
     when we wish to explain a shift item, as opposed to a reduce item. *)
  assert (0 <= i && i < n);
  {
    prefix = [];
    focus = TEmpty;
    suffix = MArray.segment_to_list rhs i n;
    comment = ""
  }

let build i rhs forest comment =
  let n = Array.length rhs in
  assert (0 <= i && i < n);
  let symbol = rhs.(i) in
  match symbol with
  | Symbol.T _ ->
      assert false
  | Symbol.N _ ->
      {
        prefix = [];
        focus = TRooted (symbol, forest);
        suffix = MArray.segment_to_list rhs (i + 1) n;
        comment = (match comment with None -> "" | Some comment -> comment)
      }

let prepend symbol forest =
  { forest with prefix = symbol :: forest.prefix }

(* -------------------------------------------------------------------------- *)

(* Display. *)

let buffer =
  Buffer.create 32768

let rec print_blank k =
  if k > 0 then begin
    Buffer.add_char buffer ' ';
    print_blank (k - 1)
  end

let print_symbol symbol =
  let word = Symbol.print false symbol in
  Buffer.add_string buffer word;
  Buffer.add_char buffer ' ';
  String.length word + 1

let print_symbols symbols =
  List.fold_left (fun offset symbol ->
    offset + print_symbol symbol
  ) 0 symbols

let print_level print_focus_root print_focus_remainder offset forest =
  print_blank offset;
  let offset = offset + print_symbols forest.prefix in
  print_focus_root forest.focus;
  let (_ : int) = print_symbols forest.suffix in
  if String.length forest.comment > 0 then begin
    Buffer.add_string buffer "// ";
    Buffer.add_string buffer forest.comment
  end;
  Buffer.add_char buffer '\n';
  print_focus_remainder offset forest.focus

let print_tree_root = function
  | TEmpty ->
      Buffer.add_string buffer ". "
  | TRooted (symbol, _) ->
      let (_ : int) = print_symbol symbol in
      ()

let rec print_forest offset forest =
  print_level print_tree_root print_tree_remainder offset forest

and print_tree_remainder offset = function
  | TEmpty ->
      ()
  | TRooted (_, forest) ->
      print_forest offset forest

let print_ctree_root = function
  | CRooted (symbol, _) ->
      let (_ : int) = print_symbol symbol in
      ()

let rec print_cforest offset cforest =
  match cforest with
  | CHole ->
      print_blank offset;
      Buffer.add_string buffer "(?)\n"
  | CCons forest ->
      print_level print_ctree_root print_ctree_remainder offset forest

and print_ctree_remainder offset = function
  | CRooted (_, cforest) ->
      print_cforest offset cforest

let wrap print channel x =
  Buffer.clear buffer;
  print 0 x;
  Buffer.output_buffer channel buffer

let print =
  wrap print_forest

let printc =
  wrap print_cforest

(* -------------------------------------------------------------------------- *)

(* [punch] turns a (tree or forest) into a pair of a (tree or forest) context
   and a residual forest, where the context is chosen maximal. In other words,
   the residual forest consists of a single level; its focus is [TEmpty]. *)

let rec punch_tree tree : (ctree * forest) option =
  match tree with
  | TEmpty ->
      None
  | TRooted (symbol, forest) ->
      let forest1, forest2 = punch_forest forest in
      Some (CRooted (symbol, forest1), forest2)

and punch_forest forest : cforest * forest =
  match punch_tree forest.focus with
  | None ->
      CHole, forest
  | Some (ctree1, forest2) ->
      CCons {
        prefix = forest.prefix;
        focus = ctree1;
        suffix = forest.suffix;
        comment = forest.comment
      }, forest2

(* [fill] fills a (tree or forest) context with a forest so as to produce
   a new (tree or forest). *)

let rec fill_tree ctree1 forest2 : tree =
  match ctree1 with
  | CRooted (symbol1, cforest1) ->
      TRooted (symbol1, fill_forest cforest1 forest2)

and fill_forest cforest1 forest2 : forest =
  match cforest1 with
  | CHole ->
      forest2
  | CCons level1 ->
      {
        prefix = level1.prefix;
        focus = fill_tree level1.focus forest2;
        suffix = level1.suffix;
        comment = level1.comment
      }

(* [common] factors the maximal common (tree or forest) context out of a pair
   of a (tree or forest) context and a (tree or forest). It returns the (tree
   or forest) context as well as the residuals of the two parameters. *)

let rec common_tree ctree1 tree2 : (ctree * cforest * forest) option =
  match ctree1, tree2 with
  | CRooted _, TEmpty ->
      None
  | CRooted (symbol1, cforest1), TRooted (symbol2, forest2) ->
      if Symbol.equal symbol1 symbol2 then
        let cforest, cforest1, forest2 =
          common_forest cforest1 forest2
        in
        Some (CRooted (symbol1, cforest), cforest1, forest2)
      else
        None

and common_forest cforest1 forest2 : cforest * cforest * forest =
  match cforest1 with
  | CHole ->
      CHole, cforest1, forest2
  | CCons forest1 ->
      if MList.equal Symbol.equal forest1.prefix forest2.prefix
      && MList.equal Symbol.equal forest1.suffix forest2.suffix
      && forest1.comment = forest2.comment
      then begin
        match common_tree forest1.focus forest2.focus with
        | None ->
            CHole, cforest1, forest2
        | Some (ctree, csubforest1, subforest2) ->
            let cforest = {
              prefix = forest1.prefix;
              focus = ctree;
              suffix = forest1.suffix;
              comment = forest1.comment
            } in
            CCons cforest, csubforest1, subforest2
      end
      else
        CHole, cforest1, forest2

(* [factor] factors the maximal common forest context out of a nonempty family
   of forests. We assume that the family is represented as a map indexed by
   items, because this is convenient for the application that we have in mind,
   but this assumption is really irrelevant. *)

let factor forests =
  Item.Map.fold (fun item forest accu ->
    match accu with

    | None ->

        (* This is the first time through the loop, so [forest] is the first
           forest that we examine. Punch it, so as to produce a maximal forest
           context and a residual forest. *)
        let context, residual = punch_forest forest in
        Some (context, Item.Map.singleton item residual)

    | Some (context, residuals) ->

        (* This is another iteration through the loop. [context] and [residuals]
           are the maximal common context and the residuals of the forests
           examined so far. *)

        (* Combine the common context obtained so far with the forest at hand.
           This yields a new, smaller common context, as well as residuals for
           the previous common context and for the forest at hand. *)
        let context, contextr, forestr = common_forest context forest in

        (* The residual forests are now: (i) the residual forest [forestr];
           and (ii) the previous residual forests [residuals], each of which
           must be placed with the residual context [contextr]. *)
        let residuals =
          Item.Map.add item forestr
            (Item.Map.map (fill_forest contextr) residuals)
        in

        Some (context, residuals)

  ) forests None
  |> Option.get (* [forests] is nonempty so the result must be [Some _] *)

(* -------------------------------------------------------------------------- *)

end (* Make *)
