#lang shplait

// Start with "letrec_ubi.rhm".

// Add support for `block` plus `def` with to have multiple `def`s,
// where the right-hand side of each `def` can see all the other
// defined names.

// A `letrec` form corresponds to a `block` with one `def`, so just
// use the same representation internally. It's simplest and easiest
// to start by changing `letrecE` to support multiple names and
// multiple right-hand sides.

module test:
  check: interp(parse('block:
                         def even = (fun (n):
                                       if n == 0
                                       | 1
                                       | odd(n + -1))
                         def odd = (fun (n):
                                      if n == 0
                                      | 0
                                      | even(n + -1))
                         odd(21)'),
                mt_env)
         ~is intV(1)

  // This example show sthe contrast between the better
  // and worse options in the `letrecE` case of `interp`
  check: interp(parse('block:
                         def x = 1
                         def y = x
                         y'),
                mt_env)
         ~is intV(1)

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

type Exp
| intE(n :: Int)
| idE(s :: Symbol)
| plusE(l :: Exp,
        r :: Exp)
| multE(l :: Exp,
        r :: Exp)
| funE(n :: Symbol,
       body :: Exp)
| appE(fn :: Exp,
       arg :: Exp)
| if0E(tst :: Exp,
       thn :: Exp,
       els :: Exp)
| letrecE(ns :: Listof(Symbol),
          rhss :: Listof(Exp),
          body :: Exp)

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

type Env = Listof(Binding)

def mt_env = []
def extend_env = cons
def append_env = append

// 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
      | 'if $tst == 0
         | $thn
         | $els':
          if0E(parse(tst),
               parse(thn),
               parse(els))
      | 'block:
           def $n = $rhs
           ...
           $body':
          letrecE(map(syntax_to_symbol,syntax_to_list('[$n, ...]')),
                  map(parse, syntax_to_list('[$rhs, ...]')),
                  parse(body))
           
      | 'letrec $name = $rhs:
           $body':
          letrecE([syntax_to_symbol(name)],
                  [parse(rhs)],
                  parse(body))
      | 'let $name = $rhs:
           $body':
          appE(funE(syntax_to_symbol(name),
                    parse(body)),
               parse(rhs))               
      | '$left + $right':
          plusE(parse(left),
                parse(right))
      | '$left * $right':
          multE(parse(left),
                parse(right))
      | 'fun ($id): $body':
          funE(syntax_to_symbol(id),
               parse(body))
      | '$fn($arg)':
          appE(parse(fn),
               parse(arg))
      | '($e)':
          parse(e)
      | ~else:
          error(#'parse, "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): 9')
         ~is funE(#'x, 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 = 1 + 2:
                  y')
         ~is appE(funE(#'x,
                       idE(#'y)),
                  plusE(intE(1), intE(2)))
  check: parse('if 1 == 0
                | 2
                | 3')
         ~is if0E(intE(1), intE(2), intE(3))
  check: parse('letrec x = 1 + 2:
                  y')
         ~is letrecE([#'x],
                     [plusE(intE(1), intE(2))],
                     idE(#'y))
  check: parse('block:
                  def x = 1
                  def y = 2
                  y + x')
         ~is letrecE([#'x,#'y],
                     [intE(1),intE(2)],
                     plusE(idE(#'y), idE(#'x))) 
  check: parse('1 2')
         ~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, body): closV(n, body, env)
  | appE(fn, arg):
      match interp(fn, env)
      | closV(n, body, c_env):
          interp(body,
                 let val = interp(arg, env):
                   extend_env(bind(n, box(some(val))),
                              c_env))
      | ~else: error(#'interp, "not a function")
  | if0E(tst, thn, els):
      match interp(tst, env)
      | intV(n): if n == 0
                 | interp(thn, env)
                 | interp(els, env)
      | ~else:
          error(#'interp, "not a number")
  | letrecE(ns, rhss, body):
      let bs = map(fun(rhs): box(none()), rhss):
        let new_env = append_env(map2(bind, ns, bs),
                                 env):
          begin:
            map2(fun(b, rhs): set_box(b, some(interp(rhs, new_env))), bs, rhss)
            interp(body, new_env)
          // This version is not as good, because it evaluates
          // all right-hand sides before setting any boxes, which
          // means that later right0hand sides can't immeidately
          // refer to earlier variables:
          #//
          let vals = map(fun (rhs): interp(rhs, new_env), rhss):
            begin:
              map2(fun(b, val): set_box(b, some(val)), bs, vals)
              interp(body, new_env)

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, box(some(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): x + x'),
                mt_env)
         ~is closV(#'x, plusE(idE(#'x), idE(#'x)), mt_env)
  check: interp(parse('let x = 5: x + x'),
                mt_env)
         ~is intV(10)
  check: interp(parse('let x = 5:
                         let x = x + 1:
                           x + x'),
                mt_env)
         ~is intV(12)
  check: interp(parse('let x = 5:
                         let y = 6:
                           x'),
                mt_env)
         ~is intV(5)
  check: interp(parse('(fun (x): x + x)(8)'),
                mt_env)
         ~is intV(16)

  check: interp(parse('if 1 == 0 | 2 | 3'),
                mt_env)
         ~is intV(3)
  check: interp(parse('if 0 == 0 | 2 | 3'),
                mt_env)
         ~is intV(2)

  check: interp(parse('letrec x = 5:
                         x + x'),
                mt_env)
         ~is intV(10)

  check: interp(parse('letrec fac = (fun (x):
                                       if x == 0
                                       | 1
                                       | x * fac(x + -1)):
                         fac(5)'),
                mt_env)
         ~is intV(120)

  check: interp(parse('1(2)'), mt_env)
         ~raises "not a function"
  check: interp(parse('1 + (fun (x): x)'), mt_env)
         ~raises "not a number"
  check: interp(parse('let bad = (fun (x): x + y):
                         let y = 5:
                           bad(2)'),
                mt_env)
         ~raises "free variable"

  check: interp(parse('if (fun (x): x) == 0 | 2 | 3'),
                mt_env)
         ~raises "not a number"
  check: interp(parse('letrec x = x: x'),
                mt_env)
         ~raises "use before initialization"

// 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 lookup(n :: Symbol, env :: Env) :: Value:
  match env
  | []: error(#'lookup, "free variable: " +& n)
  | cons(b, rst_env):
      cond
      | n == bind.name(b):
          match unbox(bind.val(b))
          | none(): error(#'lookup, "use before initialization: " +& n)
          | some(v): v
      | ~else:
          lookup(n, rst_env)

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