#lang shplait

// Typed Paret with Records and Holes: Typed Paret (as in stage 2 of the
// Type Checker assignment) plus records, with the hole, @, as the field
// of a lookup.
//
//   <expr> ::= ... everything in Typed Paret ...
//            | (record (<var> <expr>) ...)
//            | (lookup <expr> <var>)
//            | (lookup <expr> @)
//
//   <type> ::= Num | Str | Bool | (<type> -> <type>)
//            | (Record (<var> <type>) ...)
//
// This module provides the types (`Type`, `TEnv`), the surface abstract
// syntax (`Surface`), the core abstract syntax (`Expr`), the field position
// of a lookup (`Lookup`, shared by both), `parse_surface` and
// `parse_surface_str` (which accept only programs with exactly one hole,
// as the field of a lookup, and records and record types without repeated
// field names), and `same_set`. There is no parser for `Expr`: a core
// program comes from your `desugar`.

import:
  open: "sexp.rhm"

type Op
| op_plus()
| op_append()
| op_num_eq()
| op_str_eq()

type Type
| t_num()
| t_str()
| t_bool()
| t_fun(arg :: Type, ret :: Type)
| t_record(fields :: Listof(TField))

type TField
| tfield(name :: Symbol, ty :: Type)

type TEnv = Mapof(Symbol, Type)

type Surface
| s_num(n :: Int)
| s_str(s :: String)
| s_bool(b :: Boolean)
| s_op(op :: Op, left :: Surface, right :: Surface)
| s_if(cnd :: Surface, thn :: Surface, els :: Surface)
| s_lam(param :: Symbol, param_type :: Type, body :: Surface)
| s_app(fn :: Surface, arg :: Surface)
| s_var(name :: Symbol)
| s_and(left :: Surface, right :: Surface)
| s_or(left :: Surface, right :: Surface)
| s_let(name :: Symbol, name_type :: Type, rhs :: Surface, body :: Surface)
| s_record(fields :: Listof(SField))
| s_lookup(rec :: Surface, field :: Lookup)

type SField
| sfield(name :: Symbol, val :: Surface)

type Lookup
| l_hole()
| l_name(name :: Symbol)

type Expr
| e_num(n :: Int)
| e_str(s :: String)
| e_bool(b :: Boolean)
| e_op(op :: Op, left :: Expr, right :: Expr)
| e_if(cnd :: Expr, thn :: Expr, els :: Expr)
| e_lam(param :: Symbol, param_type :: Type, body :: Expr)
| e_app(fn :: Expr, arg :: Expr)
| e_var(name :: Symbol)
| e_record(fields :: Listof(Field))
| e_lookup(rec :: Expr, field :: Lookup)

type Field
| field(name :: Symbol, val :: Expr)

// Do the two lists have the same elements, ignoring order and duplicates?
fun same_set(a :: Listof(Symbol), b :: Listof(Symbol)) :: Boolean:
  block:
    fun all_in(xs :: Listof(Symbol), ys :: Listof(Symbol)) :: Boolean:
      match xs
      | []: #true
      | cons(x, rst): member(x, ys) && all_in(rst, ys)
    all_in(a, b) && all_in(b, a)

// ---------------------------------------------------------------- parser

fun parse_surface_str(str :: String) :: Surface:
  parse_surface(read_str(str))

// The helpers are local, so that this module exports only the types,
// the two parse functions, and same_set.
fun parse_surface(sexp :: Sexp) :: Surface:
  block:
    // The operator keywords, which are not writable as Shplait symbol literals.
    def sym_plus = string_to_symbol("+")
    def sym_append = string_to_symbol("++")
    def sym_num_eq = string_to_symbol("num=")
    def sym_str_eq = string_to_symbol("str=")
    def sym_colon = string_to_symbol(":")
    def sym_arrow = string_to_symbol("->")
    def sym_hole = string_to_symbol("@")

    def keywords :: Listof(Symbol) = [#'lam, #'if, #'and, #'or, #'let, #'record, #'lookup, sym_plus, sym_append, sym_num_eq, sym_str_eq, #'true, #'false, sym_hole]

    fun is_keyword(sym :: Symbol) :: Boolean:
      member(sym, keywords)

    fun parse_error(msg :: String, s :: Sexp) :: Surface:
      error(#'parse, msg +& ": " +& sexp_to_text(s))

    fun parse_type(s :: Sexp) :: Type:
      cond
      | sexp_is_symbol(s) && sexp_to_symbol(s) == #'Num: t_num()
      | sexp_is_symbol(s) && sexp_to_symbol(s) == #'Str: t_str()
      | sexp_is_symbol(s) && sexp_to_symbol(s) == #'Bool: t_bool()
      | sexp_is_list(s)
          && length(sexp_to_list(s)) == 3
          && sexp_is_symbol(list_get(sexp_to_list(s), 1))
          && sexp_to_symbol(list_get(sexp_to_list(s), 1)) == sym_arrow:
          t_fun(parse_type(first(sexp_to_list(s))), parse_type(list_get(sexp_to_list(s), 2)))
      | sexp_is_list(s)
          && !is_empty(sexp_to_list(s))
          && sexp_is_symbol(first(sexp_to_list(s)))
          && sexp_to_symbol(first(sexp_to_list(s))) == #'Record:
          t_record(parse_type_fields(rest(sexp_to_list(s)), [], s))
      | ~else: error(#'parse, "not a type: " +& sexp_to_text(s))

    // The fields of (Record (<var> <type>) ...), with no name repeated.
    fun parse_type_fields(parts :: Listof(Sexp), seen :: Listof(Symbol), whole :: Sexp) :: Listof(TField):
      match parts
      | []: []
      | cons(a, rst):
          if sexp_is_list(a) && length(sexp_to_list(a)) == 2 && sexp_is_symbol(first(sexp_to_list(a)))
          | block:
              def name = sexp_to_symbol(first(sexp_to_list(a)))
              if member(name, seen)
              | error(#'parse, "repeated field name: " +& to_string(name))
              | cons(tfield(name, parse_type(list_get(sexp_to_list(a), 1))), parse_type_fields(rst, cons(name, seen), whole))
          | error(#'parse, "not a type: " +& sexp_to_text(whole))

    fun parse_param(s :: Sexp) :: Symbol:
      block:
        def sym = sexp_to_symbol(s)
        if is_keyword(sym)
        | error(#'parse, "keyword used as a parameter name: " +& to_string(sym))
        | sym

    // Is s a list of n parts whose first is a symbol and second is `:`?
    fun is_annotated(s :: Sexp, n :: Int) :: Boolean:
      sexp_is_list(s)
        && length(sexp_to_list(s)) == n
        && sexp_is_symbol(first(sexp_to_list(s)))
        && sexp_is_symbol(list_get(sexp_to_list(s), 1))
        && sexp_to_symbol(list_get(sexp_to_list(s), 1)) == sym_colon

    // Is s of the form (<var> <expr>)?
    fun is_binding(s :: Sexp) :: Boolean:
      sexp_is_list(s) && length(sexp_to_list(s)) == 2 && sexp_is_symbol(first(sexp_to_list(s)))

    fun parse_field_name(s :: Sexp) :: Lookup:
      block:
        def sym = sexp_to_symbol(s)
        cond
        | sym == sym_hole: l_hole()
        | is_keyword(sym): error(#'parse, "keyword used as a field name: " +& to_string(sym))
        | ~else: l_name(sym)

    // The fields of (record (<var> <expr>) ...), with no name repeated.
    fun parse_fields(args :: Listof(Sexp), whole :: Sexp) :: Listof(SField):
      block:
        fun go(args :: Listof(Sexp), seen :: Listof(Symbol)) :: Listof(SField):
          match args
          | []: []
          | cons(a, rst):
              if is_binding(a)
              | block:
                  def parts = sexp_to_list(a)
                  def name = sexp_to_symbol(first(parts))
                  cond
                  | is_keyword(name): error(#'parse, "keyword used as a field name: " +& to_string(name))
                  | member(name, seen): error(#'parse, "repeated field name: " +& to_string(name))
                  | ~else: cons(sfield(name, p(list_get(parts, 1))), go(rst, cons(name, seen)))
              | error(#'parse, "a record field is (name value): " +& sexp_to_text(whole))
        go(args, [])

    fun p(s :: Sexp) :: Surface:
      cond
      | sexp_is_number(s): s_num(sexp_to_number(s))
      | sexp_is_string(s): s_str(sexp_to_string(s))
      | sexp_is_symbol(s): parse_symbol(sexp_to_symbol(s))
      | sexp_is_list(s): parse_list(sexp_to_list(s), s)
      | ~else: parse_error("cannot parse", s)

    fun parse_symbol(sym :: Symbol) :: Surface:
      cond
      | sym == #'true: s_bool(#true)
      | sym == #'false: s_bool(#false)
      | sym == sym_hole: error(#'parse, "the hole, @, may appear only as the field of a lookup")
      | is_keyword(sym): error(#'parse, "keyword used as a variable: " +& to_string(sym))
      | ~else: s_var(sym)

    fun parse_list(parts :: Listof(Sexp), whole :: Sexp) :: Surface:
      cond
      | is_empty(parts): parse_error("empty parentheses", whole)
      | sexp_is_symbol(first(parts)) && is_keyword(sexp_to_symbol(first(parts))):
          parse_form(sexp_to_symbol(first(parts)), rest(parts), whole)
      | length(parts) == 2: s_app(p(first(parts)), p(list_get(parts, 1)))
      | ~else: parse_error("function application takes exactly one argument", whole)

    // The n-th argument of a form that must have exactly k, parsed.
    fun arg(args :: Listof(Sexp), k :: Int, n :: Int, whole :: Sexp) :: Surface:
      if length(args) == k
      | p(list_get(args, n))
      | parse_error("wrong number of parts", whole)

    fun parse_form(kw :: Symbol, args :: Listof(Sexp), whole :: Sexp) :: Surface:
      cond
      | kw == sym_plus: s_op(op_plus(), arg(args, 2, 0, whole), arg(args, 2, 1, whole))
      | kw == sym_append: s_op(op_append(), arg(args, 2, 0, whole), arg(args, 2, 1, whole))
      | kw == sym_num_eq: s_op(op_num_eq(), arg(args, 2, 0, whole), arg(args, 2, 1, whole))
      | kw == sym_str_eq: s_op(op_str_eq(), arg(args, 2, 0, whole), arg(args, 2, 1, whole))
      | kw == #'if: s_if(arg(args, 3, 0, whole), arg(args, 3, 1, whole), arg(args, 3, 2, whole))
      | kw == #'and: s_and(arg(args, 2, 0, whole), arg(args, 2, 1, whole))
      | kw == #'or: s_or(arg(args, 2, 0, whole), arg(args, 2, 1, whole))
      | kw == #'let:
          if length(args) == 2 && is_annotated(first(args), 4)
          | block:
              def binding = sexp_to_list(first(args))
              s_let(parse_param(first(binding)), parse_type(list_get(binding, 2)),
                    p(list_get(binding, 3)), p(list_get(args, 1)))
          | parse_error("let takes a (name : type value) binding and a body", whole)
      | kw == #'record: s_record(parse_fields(args, whole))
      | kw == #'lookup:
          if length(args) == 2 && sexp_is_symbol(list_get(args, 1))
          | s_lookup(p(first(args)), parse_field_name(list_get(args, 1)))
          | parse_error("lookup takes a record and a field name or @", whole)
      | kw == #'lam:
          if length(args) == 2 && is_annotated(first(args), 3)
          | block:
              def param = sexp_to_list(first(args))
              s_lam(parse_param(first(param)), parse_type(list_get(param, 2)), p(list_get(args, 1)))
          | parse_error("lam takes a (name : type) parameter and a body", whole)
      | ~else: parse_error("not an expression", whole)

    fun holes(s :: Surface) :: Int:
      block:
        fun all(ss :: Listof(SField)) :: Int:
          match ss
          | []: 0
          | cons(f, rst):
              match f
              | sfield(n, v): holes(v) + all(rst)
        match s
        | s_lookup(r, l):
            match l
            | l_hole(): 1 + holes(r)
            | l_name(n): holes(r)
        | s_record(fs): all(fs)
        | s_op(op, l, r): holes(l) + holes(r)
        | s_if(c, t, f): holes(c) + holes(t) + holes(f)
        | s_lam(x, xt, b): holes(b)
        | s_app(f, a): holes(f) + holes(a)
        | s_and(l, r): holes(l) + holes(r)
        | s_or(l, r): holes(l) + holes(r)
        | s_let(x, xt, rhs, b): holes(rhs) + holes(b)
        | ~else: 0

    def result = p(sexp)
    if holes(result) == 1
    | result
    | error(#'parse, "a program must contain exactly one hole, @, as the field of a lookup: " +& sexp_to_text(sexp))

module test:
  check:
    parse_surface_str("(let (rec : (Record (a Num) (b Num)) (record (a 1) (b 2))) (lookup rec @))")
    ~is s_let(#'rec, t_record([tfield(#'a, t_num()), tfield(#'b, t_num())]),
              s_record([sfield(#'a, s_num(1)), sfield(#'b, s_num(2))]),
              s_lookup(s_var(#'rec), l_hole()))
  check:
    parse_surface_str("(lam (r : (Record)) (lookup r @))")
    ~is s_lam(#'r, t_record([]), s_lookup(s_var(#'r), l_hole()))
  check:
    parse_surface_str("(lam (r : (Record (a Num) (a Str))) (lookup r @))")
    ~raises "repeated field name"
  check:
    parse_surface_str("(lam (x : Num) @)")
    ~raises "only as the field of a lookup"
  check:
    parse_surface_str("(lookup (record (a 1)) a)")
    ~raises "exactly one hole"
  check:
    same_set([#'a], [#'a, #'a])
    ~is #true