abt.sml

signature EQ =
sig
  type t
  val eq : t * t -> bool
end

signature SHOW =
sig
  type t
  val toString : t -> string
end

signature LIST_PRETTY =
sig
  val pretty : ('a -> string) -> string * 'a list -> string
end

structure ListPretty : LIST_PRETTY =
struct
  fun pretty f (sep, xs) =
    let
      fun go [] = ""
        | go (x :: []) = f x
        | go (x :: xs) = f x ^ sep ^ go xs
    in
      go xs
    end
end

signature SORT =
sig
  type sort
  include
    sig
      include EQ
      include SHOW
    end
    where type t = sort
end

signature VALENCE =
sig
  type sort
  type valence = sort list * sort
  include
    sig
      include SHOW
      include EQ
    end
    where type t = valence
end

signature ARITY =
sig
  structure Sort : SORT
  structure Valence : VALENCE
  sharing type Valence.sort = Sort.sort

  type arity = Valence.t list * Sort.t

  include
    sig
      include SHOW
      include EQ
    end
    where type t  = arity
end

functor Valence (Sort : SORT) : VALENCE =
struct
  type sort = Sort.t
  type valence = sort list * sort
  type t = valence

  fun eq ((sorts, sigma), (sorts', sigma')) =
    ListPair.allEq Sort.eq (sorts, sorts')
      andalso Sort.eq (sigma, sigma')

  fun toString ([], sigma) = Sort.toString sigma
    | toString (sorts, sigma) =
      let
        val sorts' = ListPretty.pretty Sort.toString (", ", sorts)
        val sigma' = Sort.toString sigma
      in
        "(" ^ sorts' ^ ")" ^ sigma'
      end
end

functor Arity (Sort : SORT) : ARITY =
struct
  structure Sort = Sort
  structure Valence = Valence (Sort)
  type valence = Sort.t list * Sort.t
  type arity = valence list * Sort.t
  type t = arity

  fun eq ((valences, sigma), (valences', sigma')) =
    ListPair.allEq Valence.eq (valences, valences')
      andalso Sort.eq (sigma, sigma')

  fun toString ([], sigma) = Sort.toString sigma
    | toString (valences, sigma) =
      let
        val valences' = ListPretty.pretty Valence.toString (", ", valences)
        val sigma' = Sort.toString sigma
      in
        "(" ^ valences' ^ ")" ^ sigma'
      end
end

(* An arity-indexed family of operators *)
signature OPERATOR =
sig
  structure Arity : ARITY

  type operator
  include
    sig
      include EQ
      include SHOW
    end
    where type t = operator

  val arity : operator -> Arity.arity
end

signature VARIABLE =
sig
  type variable
  val named : string -> variable

  include
    sig
      include EQ
      include SHOW
    end
    where type t = variable

  val compare : variable * variable -> order
  val clone : variable -> variable
end

signature ABT =
sig
  structure Variable : VARIABLE
  structure Operator : OPERATOR

  type variable = Variable.t
  type operator = Operator.t
  type sort = Operator.Arity.Sort.t
  type 'a spine = 'a list

  type abt
  include EQ where type t = abt

  datatype 'a view =
      ` of variable
    | $ of operator * 'a spine
    | \ of variable * 'a

  val map : ('a -> 'b) -> 'a view -> 'b view
  val check : abt view * Operator.Arity.Valence.t -> abt
  val infer : abt -> Operator.Arity.Valence.t * abt view
end

signature ABT_UTIL =
sig
  include ABT

  datatype star = STAR of star view | EMB of abt
  val checkStar : star * Operator.Arity.Valence.t -> abt
end

functor ShowAbt (Abt : ABT) :> SHOW where type t = Abt.abt =
struct
  open Abt infix $ infixr \
  type t = abt

  fun toString e =
    case #2 (infer e) of
         `x => Variable.toString x
       | x \ e => Variable.toString x ^ "." ^ toString e
       | theta $ es =>
           Operator.toString theta ^ "(" ^ ListPretty.pretty toString ("; ", es) ^ ")"
end

functor AbtUtil (Abt : ABT) : ABT_UTIL =
struct
  open Abt

  datatype star = STAR of star view | EMB of abt

  infixr 5 \
  infix 5 $

  fun checkStar (STAR (`x) , valence) = check (`x, valence)
    | checkStar (STAR (x \ ast), valence as (sigma :: sorts, tau)) =
      let
        val e = checkStar (ast, (sorts, tau))
      in
        check (x \ e, valence)
      end
    | checkStar (STAR (theta $ asts), valence) =
      let
        val (valences, _) = Operator.arity theta
        val es = ListPair.mapEq checkStar (asts, valences)
      in
        check (theta $ es, valence)
      end
    | checkStar (EMB e, valence) =
      let
        val (valence', e') = infer e
        val true = Operator.Arity.Valence.eq (valence, valence')
      in
        e
      end
    | checkStar _ = raise Match

end

functor Abt (structure Variable : VARIABLE and Operator : OPERATOR) : ABT =
struct

  structure Variable = Variable and Operator = Operator and Arity = Operator.Arity

  type variable = Variable.t
  type operator = Operator.t
  type sort = Arity.Sort.t
  type 'a spine = 'a list

  datatype abt =
      FV of variable * sort
    | BV of int * sort
    | ABS of variable * sort * abt
    | APP of operator * abt spine

  datatype 'a view =
      ` of variable
    | $ of operator * 'a spine
    | \ of variable * 'a

  infixr 5 \
  infix 5 $

  fun map f e =
    case e of
        ` x => ` x
     | x \ e => x \ f e
     | theta $ es => theta $ List.map f es

  fun shiftVariable v n e =
    case e of
        FV (v', sigma) =>
        if Variable.eq (v, v') then
          BV (n, sigma)
        else
          FV (v', sigma)
      | BV v => BV v
      | ABS (x, sigma, e') => ABS (x, sigma, shiftVariable v (n + 1) e')
      | APP (theta, es) => APP (theta, List.map (shiftVariable v n) es)

  fun addVariable v n e =
    case e of
        FV v' => FV v'
      | BV (v' as (m, sigma)) =>
        if m = n then FV (v, sigma) else BV v'
      | ABS (x, sigma, e) => ABS (x, sigma, addVariable v (n + 1) e)
      | APP (theta, es) => APP (theta, List.map (addVariable v n) es)

  fun check (` x, ([], sigma)) = FV (x, sigma)
    | check (x \ e, (sigma::sorts, tau)) = ABS (x, sigma, shiftVariable x 0 e)
    | check (theta $ es, ([], tau)) =
      let
        val (valences, tau') = Operator.arity theta
        val true = Arity.Sort.eq (tau, tau')
        fun chkInf (valence : Arity.Valence.t, e) =
          let
            val (valence' : Arity.Valence.t, _) = infer e
          in
            if Arity.Valence.eq (valence, valence') then
              e
            else
              raise Fail "valence mismatch"
          end

        val es' = ListPair.mapEq chkInf (valences, es)
      in
        APP (theta, es')
      end
    | check _ = raise Match

  and infer (FV (v, sigma)) = (([], sigma), ` v)
    | infer (BV i) = raise Fail "Impossible: Unexpected bound variable"
    | infer (ABS (x, sigma, e)) =
      let
        val x' = Variable.clone x
        val ((sorts, tau), e') = infer e
        val valence = (sigma :: sorts, tau)
      in
        (valence, x' \ addVariable x' 0 e)
      end
    | infer (APP (theta, es)) =
      let
        val (_, tau) = Operator.arity theta
      in
        (([], tau), theta $ es)
      end

  structure Eq : EQ =
  struct
    type t = abt
    fun eq (FV (v, _), FV (v', _)) = Variable.eq (v, v')
      | eq (BV (i, _), BV (j, _)) = i = j
      | eq (ABS (_, _, e), ABS (_, _, e')) = eq (e, e')
      | eq (APP (theta, es), APP (theta', es')) =
          Operator.eq (theta, theta') andalso ListPair.allEq eq (es, es')
      | eq _ = false
  end

  open Eq
end

functor Variable () :> VARIABLE =
struct
  type variable = int * string
  type t = variable
  val counter = ref 0
  fun named a =
    let
      val i = !counter
      val _ = counter := i + 1
    in
      (i, a)
    end

  fun compare ((i, _), (j, _)) =
    Int.compare (i, j)

  fun eq ((i : int, _), (j, _)) =
    i = j

  fun clone (_, a) =
    named a

  fun toString (_, a) = a
end

structure Example =
struct
  structure V = Variable ()
  structure O =
  struct
    structure Sort =
    struct
      datatype sort = EXP | VAL | NAT
      type t = sort
      val eq : sort * sort -> bool = op=
      fun toString EXP = "exp"
        | toString VAL = "val"
        | toString NAT = "nat"
    end

    structure Arity = Arity (Sort)

    datatype operator = LAM | AP | NUM | ZE | SU | RET
    type t = operator

    fun eq (x:t, y) = x = y
    fun toString LAM = "lam"
      | toString AP = "ap"
      | toString NUM = "num"
      | toString ZE = "ze"
      | toString SU = "su"
      | toString RET = "ret"

    fun ->> (rho, tau) = (rho, tau)
    infixr ->>

    fun c x = [] ->> x

    local
      open Sort
    in
      fun arity LAM = [[EXP] ->> EXP] ->> VAL
        | arity RET = [c VAL] ->> EXP
        | arity AP = [c EXP, c EXP] ->> EXP
        | arity NUM = [c NAT] ->> VAL
        | arity ZE = c NAT
        | arity SU = [c NAT] ->> NAT
    end
  end

  structure Abt = AbtUtil(Abt (structure Operator = O and Variable = V))
  open O O.Sort Abt

  infixr 5 \
  infix 5 $

  val $$ = STAR o op$
  infix 5 $$

  val \\ = STAR o op\
  infixr 4 \\

  val `` = STAR o `
  val f = V.named "f"
  val example = checkStar (LAM $$ [f \\ AP $$ [``f, RET $$ [NUM $$ [SU $$ [ZE $$ []]]]]], c VAL)
end