#lang shplait

// Task: define and use `distance` in Moe
//
// We can write this as a new `interp` example.
//
// The `distance` function takes a record with `x`
// and `y` fields and returns the sum of those fields.
module test:
  check: interp(parse('let distance = (fun (r): r.x + r.y):
                         distance({x: 5, y:3})'),
                mt_env)
         ~is intV(8)

// Task: define `swap` in Moe
//
// The `swap` function takes a record with `x` and `y`
// fields and swaps the values of those fields
module test:
  def swap1 = '(fun (r): {x: r.y, y: r.x})'
  check: interp(parse('let swap = $swap1:
                         swap({x: 10, y:1})'),
                mt_env)
         ~is recV([#'x,#'y],[intV(1),intV(10)])
  check: interp(parse('let swap = $swap1:
                         swap({x: 10, y:1, z: 4})'),
                mt_env)
         ~is recV([#'x,#'y],[intV(1),intV(10)])

  def swap2 = '(fun (r): (r with (x = r.y)) with (y = r.x))'
  check: interp(parse('let swap = $swap2:
                         swap({x: 10, y:1, z: 4})'),
                mt_env)
         ~is recV([#'x,#'y, #'z],[intV(1),intV(10), intV(4)])
  check: interp(parse('let swap = $swap2:
                         swap({x: 10, y:1})'),
                mt_env)
         ~is recV([#'x,#'y],[intV(1),intV(10)])

// Task: Change Moe to have an form for adding a field:
//   <expr> adding (<id>     =     <expr>)
//     ^--record     ^--new field   ^---new field value
// which creates a new record like the one from the
// first expression, but with an additional field.

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

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)
| recordE(s :: Listof(Symbol),
          args :: Listof(Exp))
| getE(rec :: Exp,
       n :: Symbol)
| setE(rec :: Exp,
       n :: Symbol,
       val :: Exp)
| addingE(rec :: Exp,
          newField :: Symbol,
          val :: Exp)
type Binding
| bind(name :: Symbol,
       val :: Value)

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
      | '{ $name: $field_exp, ...}':
          recordE(map(syntax_to_symbol, syntax_to_list('[$name, ...]')),
                  map(parse, syntax_to_list('[$field_exp, ...]')))
      | '$exp with ($name = $field_exp)':
          setE(parse(exp), syntax_to_symbol(name), parse(field_exp))
      | '$exp adding ($name = $fieldExp)':
          addingE(parse(exp), syntax_to_symbol(name), parse(fieldExp))
          
      | 'let $name = $rhs:
           $body':
          letE(syntax_to_symbol(name),
               parse(rhs),
               parse(body))
      | '$left + $right':
          plusE(parse(left),
                parse(right))
      | '$left * $right':
          multE(parse(left),
                parse(right))
      | '$exp . $name':
          getE(parse(exp), syntax_to_symbol(name))
      | '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('1 adding (x = 2)')
         ~is addingE(intE(1), #'x, intE(2))
  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('{ x: 2, y: 3 }')
         ~is recordE([#'x, #'y],
                     [intE(2), intE(3)])
  check: parse('(1 + 2).a')
         ~is getE(plusE(intE(1), intE(2)), #'a)
  check: parse('(1 + 2) with (a = 7)')
         ~is setE(plusE(intE(1), intE(2)), #'a, intE(7))

  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))
  | letE(n, rhs, body):
      interp(body,
             extend_env(bind(n, interp(rhs, env)),
                        env))
  | funE(n, body): closV(n, body, env)
  | addingE(rec, name, fieldExp) :
      match interp(rec, env)
      | recV(names, vals) :
          if member(name, names)
          | error(#'interp, "already a field")
          | recV(append(names, cons(name,[])),
                 append(vals, cons(interp(fieldExp,env),[])))
      | ~else: error(#'interp, "not a record")
  | 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")
  | recordE(ns, as):
      recV(ns,
           map(fun (a): interp(a, env),
               as))
  | getE(a, n):
      match interp(a, env)
      | recV(ns, vs): find(n, ns, vs)
      | ~else: error(#'interp, "not a record")
  | setE(a, n, v):
      match interp(a, env)
      | recV(ns, vs):
          recV(ns, update(n, interp(v, env), ns, vs))
      | ~else: error(#'interp, "not a record")

module test:
  check: interp(parse('{x : 2} adding (y = 3)'), mt_env)
         ~is recV([#'x,#'y],[intV(2),intV(3)])
  check: interp(parse('1 adding (y = 3)'), mt_env)
         ~raises "not a record"
  check: interp(parse('{x : 2} adding (x = 3)'), mt_env)
         ~raises "already a field"
  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): 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('{ a: 1 + 1,
                         b: 2 + 2 }'),
                mt_env)
         ~is recV([#'a, #'b],
                  [intV(2), intV(4)])
  check: interp(parse('{ a: 1 + 1,
                         b: 2 + 2 }.a'),
                mt_env)
         ~is intV(2)
  check: interp(parse('{ a: 1 + 1,
                         b: 2 + 2 }.b'),
                mt_env)
         ~is intV(4)
  check: interp(parse('{ a: 1 + 1,
                         b: 2 + 2 } with (a = 5)'),
                mt_env)
         ~is recV([#'a, #'b],
                  [intV(5), intV(4)])
  check: interp(parse('let r1 = { a: 1 + 1,
                                  b: 2 + 2 }:
                         let r2 = (r1 with (a = 5)):
                           r1.a + r2.a'),
                mt_env)
         ~is intV(7)

  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('6 . x'), mt_env)
         ~raises "not a record"
  check: interp(parse('6 with (x = 9)'), mt_env)
         ~raises "not a record"
               
// 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):
          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, 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)

// find & update ----------------------------------------

// Takes a name and two parallel lists, returning an item from the
// second list where the name matches the item from the first list.
fun find(n :: Symbol,
         ns :: Listof(Symbol),
         vs :: Listof(Value)) :: Value:
  match ns
  | []: error(#'interp, "no such field: " +& n)
  | cons(ns_n, ns_rst): if n == ns_n
                        | first(vs)
                        | find(n, ns_rst, rest(vs))

// Takes a name n, value v, and two parallel lists, returning a list
// like the second of the given lists, but with v in place
// where n matches the item from the first list.
fun update(n :: Symbol,
           v :: Value,
           ns :: Listof(Symbol),
           vs :: Listof(Value)) :: Listof(Value):
  match ns
  | []: error(#'interp, "no such field: " +& n)
  | cons(ns_n, ns_rst): if n == ns_n
                        | cons(v, rest(vs))
                        | cons(first(vs), update(n, v, ns_rst, rest(vs)))

module test:
  check: find(#'a, [#'a, #'b], [intV(1), intV(2)])
         ~is intV(1)
  check: find(#'b, [#'a, #'b], [intV(1), intV(2)])
         ~is intV(2)
  check: find(#'a, [], [])
         ~raises "no such field"

  check: update(#'a, intV(0), [#'a, #'b], [intV(1), intV(2)])
         ~is [intV(0), intV(2)]
  check: update(#'b, intV(0), [#'a, #'b], [intV(1), intV(2)])
         ~is [intV(1), intV(0)]
  check: update(#'a, intV(0), [], [])
         ~raises "no such field"