lamumu.ml

(* https://dl.acm.org/doi/10.1145/3720507 *)

module Syntax = struct
  type id = string
  type producer_id = id
  type consumer_id = id
  type constructor_id = id
  type deconstructor_id = id
  type type_id = id

  type base_type = Integer | User of type_id
  type ty =
    | Prd of base_type
    | Cns of base_type

  type context = (id * ty) list

  type expression =
    | Producer of producer
    | Consumer of consumer
  and producer =
    | Variable of producer_id
    | Mu of consumer_id * statement
    | Construction of constructor_id * expression list
    | Cocase of (deconstructor_id * context * statement) list
    | Integer of int
  and consumer =
    | Variable of consumer_id
    | Mu_tilde of producer_id * statement
    | Deconstruction of deconstructor_id * expression list
    | Case of (constructor_id * context * statement) list
  and statement = Cut of producer * consumer
  and type_meaning =
    | Data of (constructor_id * context) list
    | Codata of (deconstructor_id * context) list
  and type_declaration = type_id * type_meaning
  and program = (type_declaration list) * statement
end

module Typing = struct
  module S = Syntax

  type context = (S.type_declaration list) * S.context

  exception Bad of string

  let error s = raise (Bad s)

  let find_or_error (f : 'a -> bool) (l : 'a list) (s : string) : 'a =
    match List.find_opt f l with
      | Some t -> t
      | None -> error s

  let lookup_variable (v : S.id) ((_, vs) : context) : S.ty =
    snd (find_or_error (fun (v', _) -> v = v')
                       vs
                       "unbound variable")

  let add_typing (v : S.id) (t : S.ty) ((types, vs) : context) : context =
    (types, (v, t) :: vs)

  let find_from_constructors (constructors : S.constructor_id list)
                             ((types, _) : context)
                             : S.type_declaration =
     find_or_error (fun (type_id, meaning) -> match meaning with
                      | S.Data constructors' ->
                          constructors' |> List.for_all (fun (c, _) ->
                            constructors |> List.exists (fun c' -> c = c'))
                      | S.Codata _ -> false)
                   types
                   "couldn't find datatype"

  let find_from_deconstructors (deconstructors : S.deconstructor_id list)
                               ((types, _) : context)
                               : S.type_declaration =
    find_or_error (fun (type_id, meaning) -> match meaning with
                     | S.Data _ -> false
                     | S.Codata deconstructors' ->
                         deconstructors' |> List.for_all (fun (c, _) ->
                           deconstructors |> List.exists (fun c' -> c = c')))
                  types
                  "couldn't find datatype"

  let fst3 (a, b, c) = a

  let rec infer_producer (p : S.producer) (context : context) : S.ty = match p with
    | Variable x -> lookup_variable x context
    | Mu (a, s) -> error "can't infer mu"
    | Construction (c, es) ->
        let (type_id, meaning) = find_from_constructors [c] context in
        (match meaning with
           | Data constructors ->
               let context' = List.assoc c constructors in
               check_expressions es context' context;
               Prd (User type_id)
           | _ -> assert false)
    | Cocase branches ->
        let (type_id, meaning) = find_from_deconstructors (List.map fst3 branches) context in
        (match meaning with
           | Codata deconstructors ->
               List.iter (fun (d, context', s) ->
                            let (types, typings) = context in
                            let new_context = (types, context' @ typings) in
                            check_statement s new_context)
                         branches;
               Prd (User type_id)
           | _ -> assert false)
    | Integer _ -> Prd Integer

  and infer_consumer (c : S.consumer) (context : context) : S.ty = match c with
    | Variable a -> lookup_variable a context
    | Mu_tilde (x, s) -> error "can't infer mutilde"
    | Deconstruction (d, es) ->
        let (type_id, meaning) = find_from_deconstructors [d] context in
        (match meaning with
           | Codata deconstructors ->
               let context' = List.assoc d deconstructors in
               check_expressions es context' context;
               Cns (User type_id)
           | _ -> assert false)
    | Case branches ->
        let (type_id, meaning) = find_from_constructors (List.map fst3 branches) context in
        (match meaning with
           | Data constructors ->
               List.iter (fun (c, context', s) ->
                            let (types, typings) = context in
                            let new_context = (types, context' @ typings) in
                            check_statement s new_context)
                         branches;
               Cns (User type_id)
           | _ -> assert false)

  and check_producer (p : S.producer) (context : context) (t : S.ty) : unit = match p with
    | Mu (a, s) -> check_statement s (add_typing a t context)
    | _ ->
        if (infer_producer p context) <> t then
          error "type mismatch"

  and check_consumer (c : S.consumer) (context : context) (t : S.ty) : unit = match c with
    | Mu_tilde (x, s) -> check_statement s (add_typing x t context)
    | _ ->
        if (infer_consumer c context) <> t then
          error "type mismatch"

  and check_statement (s : S.statement) (context : context) : unit = match s with
    | Cut (p, c) ->
        let _ = infer_producer p context in
        let _ = infer_consumer c context in
          ()

  and check_expressions (es : S.expression list) (against : S.context) (context : context) : unit =
    List.iter2 (fun e (_, t) -> match e, t with
                  | (S.Producer p, S.Prd pt) -> check_producer p context (S.Prd pt)
                  | (S.Consumer c, S.Cns ct) -> check_consumer c context (S.Cns ct)
                  | _ -> error "mismatched expression")
               es against
end

module Evaluation = struct
  module S = Syntax
  type actual_value =
    | Data of S.constructor_id * (value list)
    | Codata of (S.deconstructor_id * S.context * S.statement) list * environment
    | Integer of int
  and value =
    | Value of actual_value
    | Covalue of (value -> value)
  and environment = (S.id * value) list

  let lookup_variable (v : S.id) (e : environment) : value = List.assoc v e
  let extend (e : environment) (v : S.id) (value : value) : environment =
    (v, value) :: e
  let extend_with_context (e : environment) (context : S.context) (values : value list) : environment =
    (List.map2 (fun (name, _) value -> (name, value)) context values) @ e

  let rec evaluate_statement (S.Cut (p, c)) (e : environment) : value =
    match (evaluate_consumer c e) with
      | Covalue k -> evaluate_producer p e k
      | _ -> assert false

  and evaluate_producer (p : S.producer) (e : environment) (k : value -> value) : value = match p with
    | Variable x -> k (lookup_variable x e)
    | Mu (a, s) -> evaluate_statement s (extend e a (Covalue k))
    | Construction (c, es) -> k (Value (Data (c, evaluate_expressions es e)))
    | Cocase branches -> k (Value (Codata (branches, e)))
    | Integer n -> k (Value (Integer n))

  and evaluate_consumer (c : S.consumer) (e : environment) : value =
    match c with
    | Variable a -> lookup_variable a e
    | Mu_tilde (x, s) -> Covalue (fun v -> evaluate_statement s (extend e x v))
    | Deconstruction (d, es) -> Covalue (function
        | Value (Codata (ds, e_new)) ->
            ds |> List.find (fun (d', _, _) -> d' = d)
            |> (fun (_, context, s) ->
                  evaluate_statement s (extend_with_context e_new context (evaluate_expressions es e)))
        | _ -> assert false)
    | Case branches -> Covalue (function
        | Value (Data (c, vs)) ->
            branches |> List.find (fun (c', _, _) -> c' = c)
            |> (fun (_, context, s) ->
                  evaluate_statement s (extend_with_context e context vs))
        | _ -> assert false)

  and evaluate_expressions (es : S.expression list) (e : environment) : value list =
      es |> List.map (function
      | S.Producer p -> evaluate_producer p e Fun.id
      | S.Consumer c -> evaluate_consumer c e)

  let evaluate_program (types, s) =
    evaluate_statement s []
end