#lang shplait

// Add
//
//    abs(EXP)
//
//  again. What is different this time?

// Add a form of conditional that doesn't require
// adding booleans:
//
//   if EXP == 0 | EXP | EXP
//
// A good name for the new `Exp` variant would be `if0E`.

// Add a new definition form that can define a new name for
// a function that is already defined:
//
//   alias SYMBOL = SYMBOL
//
// For example,
//
//   fun f(x): x + 1
//   alias g = f
//
//   g(10)
//
// should produce 11.

type Exp
| intE(n :: Int)
| idE(s :: Symbol)
| plusE(l :: Exp,
        r :: Exp)
| multE(l :: Exp,
        r :: Exp)
| appE(s :: Symbol,
       arg :: Exp)
| absE(n :: Exp)
| if0E(exp :: Exp,
       tru :: Exp,
       fal :: Exp)

type FunDef
| fd(name :: Symbol,
     arg :: Symbol,
     body :: Exp)
| alias(name :: Symbol,
        fd :: Symbol)

// An EXP is either
// - 'NUMBER'
// - 'SYMBOL'
// - 'EXP + EXP'
// - 'EXP * EXP'
// - 'SYMBOL(EXP)'
// - '(EXP)'

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 $exp == 0 | $tru | $fal':
          if0E(parse(exp), parse(tru), parse(fal))
      | '$left + $right':
          plusE(parse(left),
                parse(right))
      | '$left * $right':
          multE(parse(left),
                parse(right))
      | 'abs($arg)':
          absE(parse(arg))
      | '$sym($arg)':
          appE(syntax_to_symbol(sym),
               parse(arg))
      | '($e)':
          parse(e)
      | ~else:
          error(#'parse, "invalid input: " +& s)

fun parse_fundef(s :: Syntax) :: FunDef:
  match s
  | 'fun $name($arg): $body':
      fd(syntax_to_symbol(name),
         syntax_to_symbol(arg),
         parse(body))
  | 'alias $g = $f':
      alias(syntax_to_symbol(g), syntax_to_symbol(f))
  | ~else:
      error(#'parse, "invalid input: " +& s)

module test:
  check: parse_fundef('alias g = f')
         ~is alias(#'g, #'f)

  check: parse('if 3 == 0 | 4 | 5')
         ~is if0E(intE(3), intE(4), intE(5))

  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('double(9)')
         ~is appE(#'double, intE(9))
  check: parse('1 + double(9)')
         ~is plusE(intE(1),
                   appE(#'double, intE(9)))
  check: parse('3 * (4 + 8)')
         ~is multE(intE(3),
                   plusE(intE(4), intE(8)))
  check: parse('1 2')
         ~raises "invalid input"
  check: parse('abs(5)')
         ~is absE(intE(5))
         

  check: parse_fundef('fun double(x): x + x')
         ~is fd(#'double, #'x, plusE(idE(#'x), idE(#'x)))
  check: parse_fundef('fun f(): 1')
         ~raises "invalid input"

  def double_def = parse_fundef('fun double(x): x + x')
  def quadruple_def = parse_fundef('fun quadruple(x): double(double(x))')

// interp ----------------------------------------
fun interp(a :: Exp, defs :: Listof(FunDef)) :: Int:
  match a
  | intE(n): n
  | idE(s): error(#'interp, "free variable: " +& s)
  | plusE(l, r): interp(l, defs) + interp(r, defs)
  | multE(l, r): interp(l, defs) * interp(r, defs)
  | appE(s, arg):
      block:
        def a_def = get_fundef(s, defs)
        interp(subst(intE(interp(arg, defs)),
                     fd.arg(a_def),
                     fd.body(a_def)),
               defs)
  | absE(arg):
      block:
        def abs_def = interp(arg, defs)
        if abs_def < 0
        | - abs_def
        | abs_def
  | if0E(exp, tru, fal):
      if interp(exp, defs) == 0 | interp(tru, defs) | interp(fal, defs)
        
module test:
  check: interp(parse('if 5 == 0 | 2 | 10'), [])
         ~is 10
  check: interp(parse('if 0 == 0 | 1 | 2'), [])
         ~is 1
  check: interp(parse('double(abs(-5))'), [double_def])
         ~is 10
  check: interp(parse('triple(-5)'), [parse_fundef('fun triple(x): abs(x)')])
         ~is 5
  check: interp(parse('abs(8)'), [])
         ~is 8
  check: interp(parse('abs(-2)'), [])
         ~is 2
  check: interp(parse('abs(double(3))'), [double_def])
         ~is 6
  check: interp(parse('2'), [])
         ~is 2
  check: interp(parse('x'), [])
         ~raises "free variable"
  check: interp(parse('2 + 1'), [])
         ~is 3
  check: interp(parse('2 * 1'), [])
         ~is 2
  check: interp(parse('(2 * 3) + (5 + 8)'), [])
         ~is 19
  check: interp(parse('double(8)'),
                [double_def])
         ~is 16
  check: interp(parse('quadruple(8)'),
                [double_def, quadruple_def])
         ~is 32

// get_fundef ----------------------------------------
fun get_fundef(s :: Symbol, all_defs :: Listof(FunDef)) :: FunDef:
  get_fundef_in(s, all_defs, all_defs)

// extended to keep `all_defs` so an alias can start searching from the beginning
// of the list again
fun get_fundef_in(s :: Symbol, defs :: Listof(FunDef), all_defs :: Listof(FunDef)) :: FunDef:
  match defs
  | []: error(#'get_fundef, "undefined function: " +& s)
  | cons(a_def, rst_defs):
      match a_def
      | fd(name, arg, body) :
          if s == name
          | a_def
          | get_fundef_in(s, rst_defs, all_defs)
      | alias(name, other_name):
          if s == name
          | get_fundef_in(other_name, all_defs, all_defs)
          | get_fundef_in(s, rst_defs, all_defs)

module test:
  check: get_fundef(#'double, [double_def])
         ~is double_def
  check: get_fundef(#'double, [double_def, quadruple_def])
         ~is double_def
  check: get_fundef(#'double, [quadruple_def, double_def])
         ~is double_def
  check: get_fundef(#'quadruple, [quadruple_def, double_def])
         ~is quadruple_def
  check: get_fundef(#'double, [])
         ~raises "undefined function"
  check: get_fundef(#'also_double, [alias(#'also_double, #'double),
                                    double_def])
         ~is double_def
  
// subst ----------------------------------------
fun subst(what :: Exp, for :: Symbol, in :: Exp):
  match in
  | intE(n): in
  | idE(s): (if for == s | what | in)
  | plusE(l, r): plusE(subst(what, for, l),
                       subst(what, for, r))
  | multE(l, r): multE(subst(what, for, l),
                       subst(what, for, r))
  | appE(s, arg): appE(s, subst(what, for, arg))
  | absE(n): absE(subst(what, for, n))
  | if0E(exp, tru, fal): if0E(subst(what, for, exp), subst(what, for, tru), subst(what, for, fal))

module test:
  check: subst(intE(1), #'x, if0E(idE(#'x), intE(1), intE(2)))
         ~is if0E(intE(1), intE(1), intE(2))
  check: subst(intE(1), #'x, if0E(idE(#'x), idE(#'x), idE(#'x)))
         ~is if0E(intE(1), intE(1), intE(1))
  check: subst(intE(7), #'x, absE(idE(#'x)))
         ~is absE(intE(7))
  check: subst(intE(8), #'x, intE(9))
         ~is intE(9)
  check: subst(intE(8), #'x, idE(#'x))
         ~is intE(8)
  check: subst(intE(8), #'x, idE(#'y))
         ~is idE(#'y)
  check: subst(intE(8), #'x, plusE(idE(#'x), idE(#'y)))
         ~is plusE(intE(8), idE(#'y))
  check: subst(intE(8), #'x, multE(idE(#'y), idE(#'x)))
         ~is multE(idE(#'y), intE(8))
  check: subst(intE(8), #'x, appE(#'double, idE(#'x)))
         ~is appE(#'double, intE(8))

module test:  
  check: subst(parse('8'), #'x, parse('9'))
         ~is parse('9')
  check: subst(parse('8'), #'x, parse('x'))
         ~is parse('8')
  check: subst(parse('8'), #'x, parse('y'))
         ~is idE(#'y)
  check: subst(parse('8'), #'x, parse('x + y'))
         ~is plusE(parse('8'), idE(#'y))
  check: subst(parse('8'), #'x, parse('y * x'))
         ~is multE(idE(#'y), parse('8'))
  check: subst(parse('8'), #'x, parse('double(x)'))
         ~is parse('double(8)')