#lang shplait

// Start with "poly_lambda.rhm"

// Task:
//
//   Add the usual `box`, `unbox`, and `set_box` forms
//    - using Shplait boxes
//    - with sound type checking

// Task:
// 
//  Define polymorphic `swap` in Moe, which swaps
//  box content and then returns 0:
//    swap :: forall (?a): Boxof(?a) -> Boxof(?a) -> Int

module test:
  def swap:
    '(FUN [?a]:
        (fun (b1 :: Boxof(?a)) :
           fun (b2 :: Boxof(?a)) :
             let temp :: ?a = unbox(b1):
               let dummy :: Int = set_box(b1, unbox(b2)):
                 set_box(b2, temp)))'
  def prog:
    ' let swap :: (forall(?a): Boxof(?a) -> Boxof(?a) -> Int) = $swap:
        let boxA :: Boxof(Int) = box(1):
          let boxB :: Boxof(Int) = box(2):
            let dummy :: Int = swap[Int](boxA)(boxB):
              unbox(boxA)'

  def prog2:
    ' let swap :: (forall(?a): Boxof(?a) -> Boxof(?a) -> Int) = $swap:
        let boxA :: Boxof(Int -> Int) = box(fun (x :: Int): x):
          let boxB :: Boxof(Int -> Int) = box(fun (y :: Int): y + 1):
            let dummy :: Int = swap[Int -> Int](boxA)(boxB):
              unbox(boxA)(1)'
  check: typecheck(parse(prog), mt_env)
         ~is intT()
  check: interp(parse(prog), mt_env)
         ~is intV(2)
  check: typecheck(parse(prog2), mt_env)
         ~is intT()
  check: interp(parse(prog2), mt_env)
         ~is intV(2)

type Value
| intV(n :: Int)
| closV(arg :: Symbol,
        body :: Exp,
        env :: Env)
| polyV(body :: Exp,
        env :: Env)
| boxV(b :: Boxof(Value))

type Exp
| intE(n :: Int)
| idE(s :: Symbol)
| plusE(l :: Exp,
        r :: Exp)
| multE(l :: Exp,
        r :: Exp)
| funE(n :: Symbol,
       arg_type :: Type,
       body :: Exp)
| appE(fn :: Exp,
       arg :: Exp)
| tyfunE(n :: Symbol,
         body :: Exp)
| tyappE(tyfn :: Exp,
         tyarg :: Type)
| boxE(e :: Exp)
| unboxE(b :: Exp)
| setboxE(b :: Exp, e :: Exp)

type Type
| intT()
| boolT()
| arrowT(arg :: Type,
         result :: Type)
| idT(n :: Symbol)
| forallT(n :: Symbol,
          body :: Type)
| boxofT(is :: Type)

type Binding
| bind(name :: Symbol,
       val :: Value)

type Env = Listof(Binding)

type TypeBinding
| tbind(name :: Symbol,
        ty :: Type)
| tid(name :: Symbol)

type TypeEnv = Listof(TypeBinding)

def mt_env = []
def extend_env = cons

// parse ----------------------------------------
fun parse(s :: Syntax) :: Exp:
  cond
  | syntax_is_integer(s):
      intE(syntax_to_integer(s))
  | syntax_is_symbol(s):
      idE(syntax_to_symbol(s))
  | ~else:
      match s
      | 'let $name :: $ty = $rhs:
           $body':
          appE(funE(syntax_to_symbol(name),
                    parse_type(ty),
                    parse(body)),
               parse(rhs))
      | '$left + $right':
          plusE(parse(left),
                parse(right))
      | '$left * $right':
          multE(parse(left),
                parse(right))
      | 'fun ($id :: $ty): $body':
          funE(syntax_to_symbol(id),
               parse_type(ty),
               parse(body))
      | 'FUN [? $id]: $body':
          tyfunE(syntax_to_symbol(id),
                 parse(body))
      | '$tyfn[$ty]':
          tyappE(parse(tyfn),
                 parse_type(ty))
      | 'box($arg)':
          boxE(parse(arg))
      | 'unbox($arg)':
          unboxE(parse(arg))
      | 'set_box($box_e, $val)':
          setboxE(parse(box_e), parse(val))
      | '$fn($arg)':
          appE(parse(fn),
               parse(arg))
      | '($e)':
          parse(e)
      | ~else:
          error(#'parse, "invalid input: " +& s)

fun parse_type(s :: Syntax) :: Type:
  match s
  | 'Int': intT()
  | 'Boolean': boolT()
  | '$arg1 -> $arg2 -> $result':
      // make `->` right-associative
      parse_type('$arg1 -> ($arg2 -> $result)')
  | '$arg -> $result':
      arrowT(parse_type(arg), parse_type(result))
  | '? $name': idT(syntax_to_symbol(name))
  | 'forall (? $name): $ty':
      forallT(syntax_to_symbol(name), parse_type(ty))
  | '($ty)': parse_type(ty)
  | 'Boxof($ty)':
      boxofT(parse_type(ty))
  | ~else: error(#'parse_type, "invalid input: " +& s)
           

module test:
  check: parse('2')
         ~is intE(2)
  check: parse('x')
         ~is idE(#'x)
  check: parse('2 + 1')
         ~is plusE(intE(2), intE (1))
  check: parse('3 * 4')
         ~is multE(intE(3), intE(4))
  check: parse('3 * 4 + 8')
         ~is plusE(multE(intE(3), intE(4)),
                   intE(8))
  check: parse('fun (x :: Int): 9')
         ~is funE(#'x, intT(), intE(9))
  check: parse('double(9)')
         ~is appE(idE(#'double), intE(9))
  check: parse('1 + double(9)')
         ~is plusE(intE(1),
                   appE(idE(#'double), intE(9)))
  check: parse('3 * (4 + 8)')
         ~is multE(intE(3),
                   plusE(intE(4), intE(8)))
  check: parse('let x :: Int = 1 + 2:
                  y')
         ~is appE(funE(#'x,
                       intT(),
                       idE(#'y)),
                  plusE(intE(1), intE(2)))
  check: parse('FUN [?a]:
                  fun (x :: ?a):
                    x')
         ~is tyfunE(#'a, funE(#'x, idT(#'a), idE(#'x)))
  check: parse('f[Int]')
         ~is tyappE(idE(#'f), intT())
  check: parse('1 2')
         ~raises "invalid input"

  check: parse('box(1)')
         ~is boxE(intE(1))
  check: parse('unbox(1)')
         ~is unboxE(intE(1))
  check: parse('set_box(1, 2)')
         ~is setboxE(intE(1), intE(2))

  check: parse_type('Int')
         ~is intT()
  check: parse_type('Boolean')
         ~is boolT()
  check: parse_type('Int -> Boolean')
         ~is arrowT(intT(), boolT())
  check: parse_type('Int -> (Boolean -> Int)')
         ~is arrowT(intT(), arrowT(boolT(), intT()))
  check: parse_type('Int -> Boolean -> Int')
         ~is arrowT(intT(), arrowT(boolT(), intT()))
  check: parse_type('?a')
         ~is idT(#'a)
  check: parse_type('forall (?a): ?a -> ?a')
         ~is forallT(#'a, arrowT(idT(#'a), idT(#'a)))
  check: parse_type('1')
         ~raises "invalid input"

// interp ----------------------------------------
fun interp(a :: Exp, env :: Env) :: Value:
  match a
  | intE(n): intV(n)
  | idE(s): lookup(s, env)
  | plusE(l, r): num_plus(interp(l, env), interp(r, env))
  | multE(l, r): num_mult(interp(l, env), interp(r, env))
  | funE(n, arg_type, body): closV(n, body, env)
  | appE(fn, arg):
      match interp(fn, env)
      | closV(n, body, c_env):
          interp(body,
                 extend_env(bind(n, interp(arg, env)),
                            c_env))
      | ~else: error(#'interp, "not a function")
  | tyfunE(n, body): polyV(body, env)
  | tyappE(tyfn, tyarg):
      match interp(tyfn, env)
      | polyV(body, p_env):
          interp(body, p_env)
      | ~else: error(#'interp, "not a polymorphic value")
  | boxE(e):
      def val = interp(e, env)
      boxV(box(val))
  | unboxE(e):
      def val = interp(e, env)
      unbox(boxV.b(val))
  | setboxE(b, v):
      def b_val = interp(b, env)
      def v_val = interp(v, env)
      set_box(boxV.b(b_val), v_val)
      intV(0)

module test:
  check: interp(parse('2'), mt_env)
         ~is intV(2)
  check: interp(parse('x'), mt_env)
         ~raises "free variable"
  check: interp(parse('x'),
                extend_env(bind(#'x, intV(9)), mt_env))
         ~is intV(9)
  check: interp(parse('2 + 1'), mt_env)
         ~is intV(3)
  check: interp(parse('2 * 1'), mt_env)
         ~is intV(2)
  check: interp(parse('(2 * 3) + (5 + 8)'), mt_env)
         ~is intV(19)
  check: interp(parse('fun (x :: Int):
                         x + x'),
                mt_env)
         ~is closV(#'x, plusE(idE(#'x), idE(#'x)), mt_env)
  check: interp(parse('let x :: Int = 5: x + x'),
                mt_env)
         ~is intV(10)
  check: interp(parse('let x :: Int = 5:
                         let x :: Int = x + 1:
                           x + x'),
                mt_env)
         ~is intV(12)
  check: interp(parse('let x :: Int = 5:
                         let y :: Int = 6:
                           x'),
                mt_env)
         ~is intV(5)
  check: interp(parse('(fun (x :: Int): x + x)(8)'),
                mt_env)
         ~is intV(16)

  check: interp(parse('FUN [?a]:
                         fun (x :: ?a):
                           x'),
                mt_env)
         ~is polyV(funE(#'x, idT(#'a), idE(#'x)), mt_env)
  check: interp(parse('(FUN [?a]:
                          fun (x :: ?a):
                            x)[Int]'),
                mt_env)
         ~is closV(#'x, idE(#'x), mt_env)
  check: interp(parse('let f :: (forall (?a): ?a -> ?a)
                         = (FUN [?a]:
                              fun (x :: ?a):
                                x):
                           f[Int](1)
                             + f[Int -> Int](fun (n :: Int):
                                                     n + 1)(2)'),
                mt_env)
         ~is intV(4)
  check: interp(parse('1(2)'), mt_env)
         ~raises "not a function"
  check: interp(parse('1 + (fun (x :: Int): x)'), mt_env)
         ~raises "not a number"
  check: interp(parse('let bad :: Int -> Int = (fun (x :: Int):
                                                        x + y):
                         let y :: Int = 5:
                           bad(2)'),
                mt_env)
         ~raises "free variable"
  check: interp(parse('1[Int]'), mt_env)
         ~raises "not a polymorphic value"

  check: interp(parse('box(1)'), mt_env)
         ~is boxV(box(intV(1)))
  check: interp(parse('unbox(box(1))'), mt_env)
         ~is intV(1)
  check: interp(parse('set_box(box(1), 2)'), mt_env)
         ~is intV(0)


// num_plus and num_mult ----------------------------------------
fun num_op(op :: (Int, Int) -> Int, l :: Value, r :: Value) :: Value:
  cond
  | l is_a intV && r is_a intV:
      intV(op(intV.n(l), intV.n(r)))
  | ~else:
      error(#'interp, "not a number")
fun num_plus(l :: Value, r :: Value) :: Value:
  num_op(fun (a, b): a+b, l, r)
fun num_mult(l :: Value, r :: Value) :: Value:
  num_op(fun (a, b): a*b, l, r)

module test:
  check: num_plus(intV(1), intV(2))
         ~is intV(3)
  check: num_mult(intV(3), intV(2))
         ~is intV(6)
  
// lookup ----------------------------------------
fun make_lookup(is_one :: ?a -> Boolean,
                get_name :: ?a -> Symbol,
                get_val :: ?a -> ?b):
  fun (n :: Symbol, env :: Listof(?a)) :: ?b:
    match env
    | []: error(#'lookup, "free variable")
    | cons(b, rst_env):
        cond
        | is_one(b) && n == get_name(b):
            get_val(b)
        | ~else:
            make_lookup(is_one, get_name, get_val)(n, rst_env)

def lookup = make_lookup(fun (b): #true,
                         bind.name,
                         bind.val)

module test:
  check: lookup(#'x, mt_env)
         ~raises "free variable"
  check: lookup(#'x, extend_env(bind(#'x, intV(8)), mt_env))
         ~is intV(8)
  check: lookup(#'x, extend_env(bind(#'x, intV(9)),
                                extend_env(bind(#'x, intV(8)),
                                           mt_env)))
         ~is intV(9)
  check: lookup(#'y, extend_env(bind(#'x, intV(9)),
                                extend_env(bind(#'y, intV(8)),
                                           mt_env)))
         ~is intV(8)


// typecheck ----------------------------------------
fun typecheck(a :: Exp, tenv :: TypeEnv) :: Type:
  match a
  | intE(n): intT()
  | idE(n): type_lookup(n, tenv)
  | plusE(l, r): typecheck_nums(l, r, tenv)
  | multE(l, r): typecheck_nums(l, r, tenv)
  | funE(n, arg_type, body):
      block:
        tvarcheck(arg_type, tenv)
        arrowT(arg_type,
               typecheck(body,
                         extend_env(tbind(n, arg_type),
                                    tenv)))
  | appE(fn, arg):
      match typecheck(fn, tenv)
      | arrowT(arg_type, result_type):
          if arg_type == typecheck(arg, tenv)
          | result_type
          | type_error(arg, to_string(arg_type))
      | ~else: type_error(fn, "function")
  | tyfunE(n, body):
      forallT(n, typecheck(body,
                           extend_env(tid(n),
                                      tenv)))
  | tyappE(tyfn, tyarg):
      block:
        tvarcheck(tyarg, tenv)
        match typecheck(tyfn, tenv):
        | forallT(n, body):
            type_subst(n, tyarg, body)
        | ~else:
            type_error(tyfn, "polymorphic value")
  | boxE(e):
      boxofT(typecheck(e, tenv))
  | unboxE(e):
      match typecheck(e, tenv)
      | boxofT(t):
          t
      | ~else: type_error(e, "box")
  | setboxE(b, v):
      match typecheck(b, tenv)
      | boxofT(t):
          if typecheck(v, tenv) == t
          | intT()
          | type_error(v, to_string(t))
      | ~else: type_error(b, "box")

/* let b = box(1):
      ....
       ....  set_box(b, fun (x): x)

      unbox(b) + 5*/

module test:
  check: typecheck(parse('box(1)'), mt_env)
         ~is boxofT(intT())
  check: typecheck(parse('unbox(box(1)) + 1'), mt_env)
         ~is intT()
  check: typecheck(parse('unbox(box(fun (x :: Int): x)) + 1'), mt_env)
         ~raises "no type"
  check: typecheck(parse('box((fun (x :: Int): x) + 1)'), mt_env)
         ~raises "no type"
  check: typecheck(parse('set_box(box(1), 2)'), mt_env)
         ~is intT()
  check: typecheck(parse('set_box(1, 2)'), mt_env)
         ~raises "no type"

fun typecheck_nums(l, r, tenv):
  match typecheck(l, tenv)
  | intT():
      match typecheck(r, tenv)
      | intT():
          intT()
      | ~else: type_error(r, "num")
  | ~else: type_error(l, "num")

fun type_error(a, msg):
  error(#'typecheck, "no type: " +& a +& " not " +& msg)

def type_lookup: make_lookup(fun (b): b is_a tbind,
                             tbind.name,
                             tbind.ty)
def type_var_lookup: make_lookup(fun (b): b is_a tid,
                                 tid.name,
                                 tid.name)

module test:
  check: typecheck(parse('10'), mt_env)
         ~is intT()
  check: typecheck(parse('10 + 17'), mt_env)
         ~is intT()
  check: typecheck(parse('10 * 17'), mt_env)
         ~is intT()
  check: typecheck(parse('fun (x :: Int): 12'), mt_env)
         ~is arrowT(intT(), intT())
  check: typecheck(parse('fun (x :: Int): fun (y :: Boolean): x'),
                   mt_env)
         ~is arrowT(intT(), arrowT(boolT(), intT()))

  check: typecheck(parse('(fun (x :: Int): 12)(1 + 17)'),
                   mt_env)
         ~is intT()

  check: typecheck(parse('let x :: Int = 4:
                            let f :: (Int -> Int) = (fun (y :: Int):
                                                             x + y):
                              f(x)'),
                   mt_env)
         ~is intT()

  check: typecheck(parse('FUN [?a]:
                            fun (x :: ?a):
                              x'),
                   mt_env)
         ~is forallT(#'a, arrowT(idT(#'a), idT(#'a)))
  check: typecheck(parse('(FUN [?a]: fun (x :: ?a): x)[Int]'),
                   mt_env)
         ~is arrowT(intT(), intT())
  check: typecheck(parse('let f :: (forall (?a): ?a -> ?a)
                            = (FUN [?a]:
                                 fun (x :: ?a):
                                   x):
                              f[Int](1)
                                + f[Int -> Int](fun (n :: Int):
                                                        n + 1)(2)'),
                   mt_env)
         ~is intT()

  check: typecheck(parse('1(2)'),
                   mt_env)
         ~raises "no type"
  check: typecheck(parse('(fun (x :: Boolean): x)(2)'),
                   mt_env)
         ~raises "no type"
  check: typecheck(parse('1 + (fun (x :: Int): x)'),
                   mt_env)
         ~raises "no type"
  check: typecheck(parse('(fun (x :: Int): x) * 1'),
                   mt_env)
         ~raises "no type"
  check: typecheck(parse('1[Int]'),
                   mt_env)
         ~raises "no type"

// tvarcheck ----------------------------------------
fun tvarcheck(ty, tenv):
  match ty
  | intT(): #void
  | boolT(): #void
  | arrowT(a, b):
      block:
        tvarcheck(a, tenv)
        tvarcheck(b, tenv)
  | idT(id):
      begin:
        type_var_lookup(id, tenv)
        #void
  | forallT(id, t):
      tvarcheck(t, extend_env(tid(id),
                              tenv))
  | boxofT(t):
      tvarcheck(t, tenv)

module test:
  check: tvarcheck(intT(), mt_env)
         ~is #void
  check: tvarcheck(boolT(), mt_env)
         ~is #void
  check: tvarcheck(arrowT(intT(), boolT()), mt_env)
         ~is #void
  check: tvarcheck(idT(#'a),
                   extend_env(tid(#'a), mt_env))
         ~is #void
  check: tvarcheck(forallT(#'a, idT(#'a)),
                   mt_env)
         ~is #void
  check: tvarcheck(idT(#'a), mt_env)
         ~raises "free variable"

// type_subst ----------------------------------------
fun type_subst(what :: Symbol, for :: Type, in :: Type):
  match in
  | intT(): intT()
  | boolT(): boolT()
  | arrowT(l, r): arrowT(type_subst(what, for, l),
                         type_subst(what, for, r))
  | idT(n): if what == n
            | for
            | idT(n)
  | forallT(n, body):
      cond
      | what == n: forallT(n, body)
      | is_free_type_var(n, for):          
          // If we want to replace `a` in
          // `forall (?b): ?b -> ?a` with `a`, the result
          // `forall (?b): ?b -> ?b` would be wrong, since the
          // `b` would get captured. We instead need to
          // produce `forall (?b1): ?b1 -> ?b`
          block:
            def new_n = gen_name(n, 1, for, body)
            def new_body = type_subst(n, idT(new_n), body)
            type_subst(what,
                       for,
                       forallT(new_n, new_body))
      | ~else: forallT(n, type_subst(what, for, body))
  | boxofT(t):
      boxofT(type_subst(what, for, t))

// Helper function for substitution: generates a name like `n` that is
// not currently used (as a free type variable) in `for` or `body`.
fun gen_name(n :: Symbol, i :: Int, for :: Type, body :: Type):
  block:
    def new_n = string_to_symbol(n +& i)
    if is_free_type_var(new_n, for) || is_free_type_var(new_n, body)
    | gen_name(n, i+1, for, body)
    | new_n

// Helper function for substutition: check whether a name is used as a
//;; free type variable in a type.
fun is_free_type_var(n :: Symbol, t :: Type):
  match t
  | intT(): #false
  | boolT(): #false
  | arrowT(l, r): is_free_type_var(n, l) || is_free_type_var(n, r)
  | idT(n_v): n_v == n
  | forallT(n_f, body): cond
                        | n == n_f: #false
                        | ~else: is_free_type_var(n, body)
  | boxofT(t):
      is_free_type_var(n, t)

module test:
  check: is_free_type_var(#'a, intT())
         ~is #false
  check: is_free_type_var(#'a, boolT())
         ~is #false
  check: is_free_type_var(#'a, arrowT(idT(#'b), idT(#'b)))
         ~is #false
  check: is_free_type_var(#'a, arrowT(idT(#'a), idT(#'b)))
         ~is #true
  check: is_free_type_var(#'a, idT(#'a))
         ~is #true
  check: is_free_type_var(#'a, idT(#'b))
         ~is #false
  check: is_free_type_var(#'a, forallT(#'a, idT(#'a)))
         ~is #false
  check: is_free_type_var(#'a, forallT(#'b, idT(#'a)))
         ~is #true
  check: is_free_type_var(#'a, forallT(#'b, idT(#'c)))
         ~is #false

  check: gen_name(#'a, 1, intT(), intT())
         ~is #'a1
  check: gen_name(#'a, 1, intT(), idT(#'a1))
         ~is #'a2
  check: gen_name(#'a, 1, idT(#'a1), intT())
         ~is #'a2
  check: gen_name(#'a, 1, arrowT(intT(), idT(#'a1)), intT())
         ~is #'a2
  check: gen_name(#'a, 1, forallT(#'a1, idT(#'a1)), intT())
         ~is #'a1
  check: gen_name(#'a, 1, forallT(#'b, idT(#'a1)), intT())
         ~is #'a2

  check: type_subst(#'a, boolT(), intT())
         ~is intT()
  check: type_subst(#'a, intT(), boolT())
         ~is boolT()
  check: type_subst(#'a, intT(), arrowT(idT(#'a), boolT()))
         ~is arrowT(intT(), boolT())
  check: type_subst(#'a, intT(), idT(#'a))
         ~is intT()
  check: type_subst(#'a, intT(), idT(#'b))
         ~is idT(#'b)
  check: type_subst(#'a, intT(), forallT(#'a, idT(#'a)))
         ~is forallT(#'a, idT(#'a))
  check: type_subst(#'a, intT(), forallT(#'b, idT(#'a)))
         ~is forallT(#'b, intT())
  check: type_subst(#'a, idT(#'b), forallT(#'b, idT(#'a)))
         ~is forallT(#'b1, idT(#'b))
  check: type_subst(#'a, idT(#'b), forallT(#'b, arrowT(idT(#'a),
                                                       arrowT(idT(#'b1),
                                                              idT(#'b)))))
         ~is forallT(#'b2, arrowT(idT(#'b),
                                  arrowT(idT(#'b1),
                                         idT(#'b2))))