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

(* This module defines a little language that allows the user to describe
   complex binding constructs, in the style of Weirich et al.'s "Binders
   unbound" (ICFP 2011). *)

(* Following alphaCaml and "Binders unbound", we distinguish "terms" versus
   "patterns". The syntax of these categories is as follows:

     t ::= ... (* products, sums, etc. *)
         | abstraction(p)

     p ::= ... (* products, sums, etc. *)
         | binder
         | inner(t)
         | outer(t)
         | rebind(p)
         | recursive(p)
         | repeat(p)

   On top of this syntax, we impose the restriction that [recursive] is
   forbidden under [recursive] and [repeat]. This restriction is dynamically
   checked, because checking it statically would be impossible or complicated,
   and would yield bad type error messages when it is violated.

   The combinator [inner(p)] can be viewed as sugar for [rebind(outer(p))].

   Our [rebind(p)] takes just one argument. Weirich et al.'s [rebind],
   which takes two arguments, is emulated by writing [p1 * rebind(p2)]. *)

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

(* Type definitions. *)

type 'p abstraction =
  'p

type 'bn binder =
  'bn

type 't inner =
  't

type 't outer =
  't

type 'p rebind =
  'p

type 'p repeat =
  'p

type 'p recursive =
  'p

(* Weirich et al.'s [bind] combinator, which takes a pattern and a term and
   constructs a term, can be defined in terms of the above combinators, as
   follows. *)

(* In natural syntax, this is [abstraction (p * inner(t))]. *)

type ('p, 't) bind =
  ('p * 't inner) abstraction

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

(* When traversing a pattern, our visitors maintain a richer environment than
   when traversing a term. If a normal term-traversal environment has type
   ['env], then an enriched pattern-traversal environment has type ['env penv]
   as defined below. *)

(* The meaning of the above combinators is defined by the manner in which they
   extend or use this enriched environment. *)

(* The [current] component is a mutable reference to an environment. This
   environment is extended when a [binder] combinator is visited. This
   mutable state is implicitly threaded from left to right. *)

(* The [outer] component is an immutable environment. It is transmitted in
   a top-down manner. *)

(* The components [at_binder] and [at_outer] respectively determine what
   behavior is desired at [binder] and [outer]. *)

type 'env penv = {
  current: 'env ref;      (* do NOT replace this with a mutable field! *)
  outer: 'env;
  at_binder: at_binder;
  at_outer: at_outer;
}

(* At [binder], two behaviors are possible:
   - [Extend] means that [binder] extends the current environment.
   - [Lookup] means that [binder] does NOT extend the current environment:
     instead, it looks up the environment, which presumably has been extended
     already. *)

and at_binder =
  | Extend
  | Lookup

(* At [outer], two behaviors are possible:
   - [Visit] means that the embedded subterm is visited.
     The [outer] component of the enriched environment is used.
   - [DoNotVisit] means that the embedded subterm is not visited.
     In short, in that case, nothing happens. *)

and at_outer =
  | Visit
  | DoNotVisit

(* We are in [Extend] and [Visit] mode at the beginning, when entering an
   [abstraction], and remain in this mode as long as we are outside both
   [recursive] and [repeat].

   [DoNotVisit] is used in the first phase of visiting a [recursive] construct,
   to find all of the bound names before visiting any embedded subterms.

   [Lookup] is used in the second phase of visiting a [recursive] construct,
   because the environment has been extended already during the first phase.
   It is used also under [repeat]. *)

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

(* A few auxiliary functions. *)

(* [abstraction] initializes an enriched environment. [current] and [outer]
   are initialized with the environment that exists outside the abstraction.
   The behavior is initially [Extend/Visit]. *)

let abstraction (env : 'env) : 'env penv =
  { current = ref env; outer = env; at_binder = Extend; at_outer = Visit }

(* [rebind] copies the [current] component of the environment into the [outer]
   component. In other words, it locally changes the meaning of [outer]. *)

let rebind (penv : 'env penv) : 'env penv =
  { penv with outer = !(penv.current) }

(* [recursive] is forbidden under [recursive] or [repeat]. If this rule is
   violated, raise an exception that carries a readable explanation. *)

let check_recursive_permitted (penv : 'env penv) : unit =
  match penv.at_binder, penv.at_outer with
  | Extend, Visit ->
      ()
  | _, DoNotVisit
  | Lookup, _ ->
      failwith "AlphaLib: \"recursive\" cannot be used under \"recursive\" or \"repeat\"."

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

(* [map] *)

class virtual ['self] map = object (self : 'self)

  method private virtual extend: 'bn1 -> 'env -> 'bn2 * 'env
  method private virtual lookup: 'bn1 -> 'env -> 'bn2

  (* [visit_abstraction] initializes an enriched environment using the
     auxiliary function [abstraction]. *)

  method private visit_abstraction: 'env 'p1 'p2 .
    ('env penv -> 'p1 -> 'p2) ->
    'env -> 'p1 abstraction -> 'p2 abstraction
  = fun visit_p env p1 ->
      visit_p (abstraction env) p1

  (* In [Extend] mode, [binder] extends the environment via a call to the
     virtual method [extend]. In [Lookup] mode, it uses the virtual method
     [lookup] instead, because (presumably) the environment has been extended
     already with this name. *)

  method private visit_binder:
    _ ->
    'env penv -> 'bn1 binder -> 'bn2 binder
  = fun _ penv x1 ->
      let env = !(penv.current) in
      match penv.at_binder with
      | Extend ->
          let x2, env = self#extend x1 env in
          penv.current := env;
          x2
      | Lookup ->
          self#lookup x1 env

  (* [inner(p)] can be viewed as sugar for [rebind(outer(p))]. Therefore, this
     code should ideally be generated. *)

  method private visit_inner: 'env 't1 't2 .
    ('env -> 't1 -> 't2) ->
    'env penv -> 't1 inner -> 't2 inner
  = fun visit_t penv t1 ->
      self#visit_rebind (self#visit_outer visit_t) penv t1

  (* In the [Visit] mode, an [outer] subterm is visited using the [outer]
     component of the enriched environment. In [DoNotVisit] mode, it must NOT
     be visited at all. This mode is used only in the discovery phase that
     takes place under a [recursive] construct. This discovery phase is
     performed only for its side effect, which is to extend the current
     environment; its result is discarded. Therefore, it is acceptable to
     return a dummy value here. Unfortunately, we must use an unsafe cast in
     order to produce a value of type ['t2]. An alternative would be to call
     [visit_t ... t1] anyway, but that would have the undesired effect of
     causing embedded terms to be visited twice when they appear under
     [recursive]. *)

  method private visit_outer: 'env 't1 't2 .
    ('env -> 't1 -> 't2) ->
    'env penv -> 't1 outer -> 't2 outer
  = fun visit_t penv t1 ->
      match penv.at_outer with
      | Visit ->
          visit_t penv.outer t1
      | DoNotVisit ->
          Obj.magic ()

  (* [visit_abstraction] updates the enriched environment using the auxiliary
     function [rebind]. *)

  method private visit_rebind: 'env 'p1 'p2 .
    ('env penv -> 'p1 -> 'p2) ->
    'env penv -> 'p1 rebind -> 'p2 rebind
  = fun visit_p penv p1 ->
      visit_p (rebind penv) p1

  (* A [recursive] is construct is visited in two successive phases. In the
     first phase, we use [Extend/DoNotVisit] mode, which means that the
     [current] component of the enriched environment is extended, but nothing
     else happens. The result of this phase is discarded. (It MUST be
     discarded, as we have used unsafe casts to build it.) In the second
     phase, we use [Lookup/Visit] mode, which means that the environment must
     not extended (it has been extended already). If a [binder] is encountered
     during this phase, the environment is looked up. If an embedded term is
     encountered during this phase, it is visited. *)

  method private visit_recursive: 'env 'p1 'p2 .
    ('env penv -> 'p1 -> 'p2) ->
    'env penv -> 'p1 recursive -> 'p2 recursive
  = fun visit_p penv p1 ->
      check_recursive_permitted penv;
      let _ = visit_p { penv with at_outer  = DoNotVisit } p1 in
              visit_p { penv with at_binder = Lookup     } p1

  (* A [repeat] construct causes a switch to [Lookup] mode. *)

  method private visit_repeat: 'env 'p1 'p2 .
    ('env penv -> 'p1 -> 'p2) ->
    'env penv -> 'p1 repeat -> 'p2 repeat
  = fun visit_p penv p1 ->
      let penv = { penv with at_binder = Lookup } in
      visit_p penv p1

  (* [bind(p, t)] is sugar for [abstraction (p * rebind(outer(t)))].
     Therefore, the following code should ideally be generated. We write it by
     hand in a slightly simpler form, where [visit_abstraction] is inlined. *)

  method private visit_bind: 'env 'p1 'p2 't1 't2 .
    ('env penv -> 'p1 -> 'p2) ->
    ('env -> 't1 -> 't2) ->
    'env -> ('p1, 't1) bind -> ('p2, 't2) bind
  = fun visit_p visit_t env (p1, t1) ->
      let penv = abstraction env in
      let p2 = visit_p penv p1 in
      let t2 = self#visit_inner visit_t penv t1 in
      p2, t2

end