#lang shplait

// Add ref cells to lazy Moe
// - ref(e) = make a cell
// - unref(e) = open a cell
//
// ref cells should be lazy, like everything else
// example: ref(1 + (fun (x): x)) = not an error

// challenge: write a test that returns a box
// with a forced value inside

type Value
| intV(n :: Int)
| closV(arg :: Symbol,
        body :: Exp,
        env :: Env)
/* | refV(exp :: Exp, env :: Env)
 *   ^^^ this is lazy but doesn't use a cache,
 *       so every `unref` will re-evaluate the exp
 *   we can fix by adding a cache directly or using a Thunk
 */
| refV(delay :: Thunk)

type Thunk
| delay(arg :: Exp,
        env :: Env,
        done :: Boxof(Optionof(Value)))

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)
| refE(arg :: Exp)
| unrefE(arg :: Exp)

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

type Env = Listof(Binding)

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
      | 'ref($arg)':
          refE(parse(arg))
      | 'unref($arg)':
          unrefE(parse(arg))
      | '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('ref(2)')
         ~is refE(intE(2))
  
  check: parse('1 2')
         ~raises "invalid input"
      
// interp ----------------------------------------
fun interp(a :: Exp, env :: Env) :: Value:
  match a
  | intE(n): intV(n)
  | idE(s): force(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,
                 extend_env(bind(n, delay(arg, env, box(none()))),
                            c_env))
      | ~else: error(#'interp, "not a function")
  | refE(arg):
      refV(delay(arg, env, box(none())))
  | unrefE(arg):
      match interp(arg, env)
      | refV(delay):
          force(delay)
      | ~else: error(#'interp, "not ref")

module test:
  check: interp(parse('let x = ref(2+2):
                         let y = unref(x):
                           y + unref(x)'),
                mt_env)
         ~is intV(8)
  check: interp(parse('ref(2)'),mt_env)
         ~is refV(delay(intE(2),
                        mt_env,
                        box(none())))
  check: interp(parse('unref(3)'),mt_env)
         ~raises "not ref"
  check: interp(parse('unref(ref(2))'),mt_env)
         ~is intV(2)
  check: interp(parse('let x=1:
                         unref(ref(x))'),mt_env)
         ~is intV(1)
         
  // challenge attempt 1
  def x_env = extend_env(bind(#'x, delay(refE(plusE(intE(2), intE(2))),
                                         mt_env,
                                         box(none()))),
                         [])
  check: interp(parse('let x = ref(2+2):
                         let y = unref(x):
                           ref(y + unref(x))'),
                mt_env)
         ~is refV(delay(plusE(idE(#'y), unrefE(idE(#'x))),
                        extend_env(bind(#'y, delay(unrefE(idE(#'x)),
                                                   x_env,
                                                   box(none()))),
                                   x_env),
                        box(none())))
  // challenge attempt 2
  check: interp(parse('let x = ref(2+2):
                         let y = unref(x):
                           unref(ref(y + unref(x)))'),
                mt_env)
         ~is intV(8)

  // challenge solution:
  // make ref, force it, AND return it
  interp(parse('let x = ref((fun (a): a)):
                  unref(x)(x)'),
         mt_env)

  // ---

  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, delay(intE(9), mt_env, box(none()))),
                           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('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('let x= 2+2:x'),mt_env)
         ~is intV(4)
  check: interp(parse('let x= 2+2:
                         let y = z+2:
                           x'),mt_env)
         ~is intV(4)
  check: interp(parse('ref(2)'),mt_env)
         ~is refV(delay(intE(2),
                        mt_env,
                        box(none())))
  check: interp(parse('unref(3)'),mt_env)
         ~raises "not ref"
  check: interp(parse('unref(ref(2))'),mt_env)
         ~is intV(2)
  

  time:
    interp(parse('let x2 = (fun (n): n + n):
                    let x4 = (fun (n): x2(x2(n))):
                      let x16 = (fun (n): x4(x4(n))):
                        let x256 = (fun (n): x16(x16(n))):
                          let x65536 = (fun (n): x256(x256(n))):
                            x65536(1)'),
           mt_env)
               
// force ----------------------------------------

fun force(t :: Thunk) :: Value:
  match t
  | delay(a, e, done):
      match unbox(done)
      | none():
          let v = interp(a, e):
            set_box(done, some(v))
            v
      | some(v):
          v

module test:
  check: force(delay(intE(8), mt_env, box(none())))
         ~is intV(8)
  check: block:
           def d = delay(intE(8), mt_env, box(none()))
           begin:
             force(d)
             force(d)
         ~is intV(8)
  check: force(delay(intE(8), mt_env, box(some(intV(9)))))
         ~is intV(9)
  check: force(delay(idE(#'x),
                     extend_env(bind(#'x, delay(intE(9), mt_env, box(none()))),
                                mt_env),
                     box(none())))
         ~is intV(9)

// 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) :: Thunk:
  match env
  | []: error(#'lookup, "free variable: " +& n)
  | cons(b, rst_env):
      cond
      | n == bind.name(b):
          bind.val(b)
      | ~else:
          lookup(n, rst_env)

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