#lang shplait

// FLUID_LET
// Add a `fluid_let <id> = <expr>: <expr>` form that:
// * changes <id> to have the value of the first `<expr>`
//   while evaluating the body <expr>,
// * returns the value of the body <expr>,
// * and finally changes <id> back to its original value.


type Location = Int

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)
| letE(n :: Symbol,
       rhs :: Exp,
       body :: Exp)
| funE(n :: Symbol,
       body :: Exp)
| appE(fn :: Exp,
       arg :: Exp)
| setE(var :: Symbol,
       val :: Exp)
| beginE(l :: Exp,
         r :: Exp)
| fluidE(n :: Symbol,
         lhs :: Exp,
         body :: Exp)

type Binding
| bind(name :: Symbol,
       location :: Location)

type Env = Listof(Binding)

def mt_env = []
def extend_env = cons

type Storage
| cell(location :: Location,
       val :: Value)

type Store = Listof(Storage)
def mt_store = []
def override_store = cons

type Result
| res(v :: Value,
      s :: Store)

// 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
      // fluid_let as syntactic sugar using temporary variable names, yay!
      // (beware: won't work for nested fluid lets)
      /*
      | 'fluid_let $name = $rhs:
           $body':
          parse('let __temp_name = $name:
                   let __result = (let $name = $rhs: $body):
                     begin:
                       $name := __temp_name
                       __result')
      */
      | 'let $name = $rhs:
           $body':
          letE(syntax_to_symbol(name),
               parse(rhs),
               parse(body))
      | 'fluid_let $name = $rhs:
          $body':
          fluidE(syntax_to_symbol(name),
                 parse(rhs),
                 parse(body))
      | '$left + $right':
          plusE(parse(left),
                parse(right))
      | '$left * $right':
          multE(parse(left),
                parse(right))
      | '$name := $val_exp':
          setE(syntax_to_symbol(name), parse(val_exp))
      | 'fun ($id): $body':
          funE(syntax_to_symbol(id),
               parse(body))
      | '$fn($arg)':
          appE(parse(fn),
               parse(arg))
      | 'begin:
           $effect_exp
           $result_exp':
          beginE(parse(effect_exp), parse(result_exp))
      | '($e)':
          parse(e)
      | ~else:
          error(#'parse, "invalid input: " +& s)

module test:
  /*check: parse('fluid_let x = 32:
                  x')
         ~is parse('let __temp_name = x:
                        let __result = (let x = 32:
                                          x):
                          begin:
                            x := __temp_name
                            __result')
  */              
  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 letE(#'x,
                  plusE(intE(1), intE(2)),
                  idE(#'y))
  check: parse('b := 0')
         ~is setE(#'b, intE(0))
  check: parse('begin:
                  1
                  2')
         ~is beginE(intE(1), intE(2))

  check: parse('1 2')
         ~raises "invalid input"

// reslet form ----------------------------------------
macro 'reslet ($v_id, $sto_id) = $call:
         $body':
  'match $call
   | res($v_id, $sto_id):
       $body'
  
// interp ----------------------------------------
fun interp(a :: Exp, env :: Env, sto :: Store) :: Result:
  match a
  | intE(n): res(intV(n), sto)
  | idE(s): res(fetch(lookup(s, env), sto),
                sto)
  | plusE(l, r): reslet (v_l, sto_l) = interp(l, env, sto):
                   reslet (v_r, sto_r) = interp(r, env, sto_l):
                     res(num_plus(v_l, v_r), sto_r)
  | multE(l, r): reslet (v_l, sto_l) = interp(l, env, sto):
                   reslet (v_r, sto_r) = interp(r, env, sto_l):
                     res(num_mult(v_l, v_r), sto_r)
  | letE(n, rhs, body):
      reslet (v_rhs, sto_rhs) = interp(rhs, env, sto):
        let l = new_loc(sto_rhs):
          interp(body,
                 extend_env(bind(n, l),
                            env),
                 override_store(cell(l, v_rhs),
                                sto_rhs))
  | fluidE(n, rhs, body):
      block:
        // save old value of n
        def loc_n = lookup(n,env)
        def prev_val = fetch(loc_n,sto)
        // interp rhs (we could use reslet here)
        def rhsV = interp(rhs,env,sto)
        match rhsV
        | res(rhs_val,rhs_store):
            block:
              // update n in the store, then interp body (could use reslet here too)
              def new_store = override_store(cell(loc_n,rhs_val),rhs_store)
              match interp(body,env,new_store)
              | res(body_val,body_store):
                  // return body_val and a store with the OLD value of n inside
                  def undo_store = override_store(cell(loc_n, prev_val), body_store)
                  res(body_val, undo_store)
  | funE(n, body): res(closV(n, body, env), sto)
  | appE(fn, arg):
      reslet (v_fn, sto_fn) = interp(fn, env, sto):
        match v_fn
        | closV(n, body, c_env):
            reslet (v_arg, sto_arg) = interp(arg, env, sto_fn):
              let l = new_loc(sto_arg):        
                interp(body,
                       extend_env(bind(n, l),
                                  c_env),
                       override_store(cell(l, v_arg),
                                      sto_arg))
        | ~else: error(#'interp, "not a function")
  | setE(var, val):
      let l = lookup(var, env):
        reslet (v_v, sto_v) = interp(val, env, sto):
          res(v_v,
              override_store(cell(l, v_v),
                             sto_v))
  | beginE(l, r):
      reslet (v_l, sto_l) = interp(l, env, sto):
        interp(r, env, sto_l)

module test:
  check: interp(parse('let x = 50:
                         (fluid_let x = 42:
                           x) + x'),
                mt_env,
                mt_store)
         ~is res(intV(50+42), override_store(cell(1, intV(50)),
                                             override_store(cell(1, intV(42)),
                                                            override_store(cell(1, intV(50)), mt_store))))

  check: res.v(interp(parse('let x = 2:
                               let f = (fun (y): x + y):
                                 fluid_let x = 20:
                                   f(2)'),
                      mt_env,
                      mt_store))
         ~is intV(22)
  check: interp(parse('2'), mt_env, mt_store)
         ~is res(intV(2), mt_store)
  check: interp(parse('x'), mt_env, mt_store)
         ~raises "free variable"
  check: interp(parse('x'),
                extend_env(bind(#'x, 1), mt_env),
                override_store(cell(1, intV(9)),
                               mt_store))
         ~is res(intV(9),
                 override_store(cell(1, intV(9)),
                                mt_store))
  check: interp(parse('2 + 1'), mt_env, mt_store)
         ~is res(intV(3), mt_store)
  check: interp(parse('2 * 1'), mt_env, mt_store)
         ~is res(intV(2), mt_store)
  check: interp(parse('(2 * 3) + (5 + 8)'), mt_env, mt_store)
         ~is res(intV(19), mt_store)
  check: interp(parse('fun (x): x + x'),
                mt_env,
                mt_store)
         ~is res(closV(#'x, plusE(idE(#'x), idE(#'x)), mt_env),
                 mt_store)
  check: interp(parse('let x = 5: x + x'),
                mt_env,
                mt_store)
         ~is res(intV(10),
                 override_store(cell(1, intV(5)),
                                mt_store))
  check: interp(parse('let x = 5:
                         let x = x + 1:
                           x + x'),
                mt_env,
                mt_store)
         ~is res(intV(12),
                 override_store(cell(2, intV(6)),
                                override_store(cell(1, intV(5)),
                                               mt_store)))
  check: interp(parse('let x = 5:
                         let y = 6:
                           x'),
                mt_env,
                mt_store)
         ~is res(intV(5),
                 override_store(cell(2, intV(6)),
                                override_store(cell(1, intV(5)),
                                               mt_store)))
  check: interp(parse('(fun (x): x + x)(8)'),
                mt_env,
                mt_store)
         ~is res(intV(16),
                 override_store(cell(1, intV(8)),
                                mt_store))

  check: interp(parse('let x = 5:
                         begin:
                           x := 6
                           x'),
                mt_env,
                mt_store)
         ~is res(intV(6),
                 override_store(cell(1, intV(6)),
                                override_store(cell(1, intV(5)),
                                               mt_store)))
  check: interp(parse('begin:
                         1
                         2'),
                mt_env,
                mt_store)
         ~is res(intV(2), mt_store)

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

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

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

// store operations ----------------------------------------

fun new_loc(sto :: Store) :: Location:
  1 + max_address(sto)

fun max_address(sto :: Store) :: Location:
  match sto
  | []: 0
  | cons(c, rst_sto): max(cell.location(c),
                          max_address(rst_sto))

fun fetch(l :: Location, sto :: Store) :: Value:
  match sto
  | []: error(#'interp, "unallocated location")
  | cons(c, rst_sto): if l == cell.location(c)
                      | cell.val(c)
                      | fetch(l, rst_sto)

module test:
  check: new_loc(mt_store)
         ~is 1
       
  check: max_address(mt_store)
         ~is 0
  check: max_address(override_store(cell(2, intV(9)),
                                    mt_store))
         ~is 2
  check: fetch(2, override_store(cell(2, intV(9)),
                                 mt_store))
         ~is intV(9)
  check: fetch(2, override_store(cell(2, intV(10)),
                                 override_store(cell(2, intV(9)),
                                                mt_store)))
         ~is intV(10)
  check: fetch(3, override_store(cell(2, intV(10)),
                                 override_store(cell(3, intV(9)),
                                                mt_store)))
         ~is intV(9)
  check: fetch(2, mt_store)
         ~raises "unallocated location"