Source file build_context.ml

open Fmlib
open Common
open Alba_core


module Algo = Gamma_algo.Make (Gamma_holes)
module Uni  = Unifier.Make (Gamma_holes)


module Stack =
struct
    let split (stack: 'a list): 'a * 'a list =
        match stack with
        | [] ->
            assert false (* Illegal call! *)
        | hd :: tl ->
            hd, tl

    let swap (sp: 'a) (stack: 'a list): 'a * 'a list =
        match stack with
        | hd :: tl ->
            hd, sp :: tl
        | _ ->
            assert false (* Illegal call! *)

    let rec pop (nargs: int) (sp: 'a) (stack: 'a list): 'a * 'a list =
        if nargs = 0 then
            sp, stack
        else
            match stack with
            | [] ->
                assert false (* Illegal call! *)
            | sp :: stack ->
                pop (nargs - 1) sp stack
end






type entry = {
    cnt0: int;
}


type t = {
    gh: Gamma_holes.t;
    sp: int;
    stack: int list;
    entry: entry;
    entries: entry list;
}


type type_in_context = int list * Term.typ * Gamma.t




let index_of_level (level: int) (bc: t): int =
    Gamma_holes.index_of_level level bc.gh


let string_of_term (term: Term.t) (bc: t): string =
    Term_printer.string_of_term term (Gamma_holes.context bc.gh)
let _ = string_of_term



let count (bc: t): int =
    Gamma_holes.count bc.gh

let count_base (bc: t): int =
    Gamma_holes.count_base bc.gh


let count_locals (bc: t): int =
    Gamma_holes.count_locals bc.gh


let count_bounds (bc: t): int =
    Gamma_holes.count_bounds bc.gh

let context (bc: t): Gamma.t =
    Gamma_holes.context bc.gh

let type_at_level (level: int) (bc: t): Term.typ =
    assert (level < count bc);
    Gamma_holes.type_at_level level bc.gh


let type_of_term (term: Term.t) (bc: t): Term.typ =
    Algo.type_of_term term bc.gh



let key_normal (term: Term.t) (bc: t): Term.t =
    Algo.key_normal term bc.gh


let is_kind (typ: Term.typ) (bc: t): bool =
    Algo.is_kind typ bc.gh



let required_type (bc: t): Term.typ =
    type_at_level bc.sp bc



let term_at_level (level: int) (bc: t): Term.t =
    Gamma_holes.expand
        (Term.Variable (index_of_level level bc))
        bc.gh



let top_term (bc: t): Term.t =
    term_at_level bc.sp bc


let unfilled_holes (term: Term.t) (bc: t): int list =
    Int_set.elements
        (Gamma_holes.unfilled_holes
            (count_base bc)
            term
            bc.gh)


let type_in_context (typ: Term.typ) (bc: t): type_in_context =
    unfilled_holes typ bc,
    typ,
    context bc


let required_type_in_context (bc: t): type_in_context =
    type_in_context (required_type bc) bc


let add_one_implicit
    (term: Term.t) (typ: Term.typ) (bc: t)
    : (Term.t * Term.typ * t) option
    =
    let open Term in
    match typ with
    | Pi (arg, res, _) when is_kind arg bc && has_variable 0 res ->
        Some (
            Appl (
                up1 term,
                Variable 0,
                Application_info.Implicit),
            res,
            {bc with gh = Gamma_holes.push_hole arg bc.gh}
        )
    | _ ->
        None






let add_implicits
    (term: Term.t) (typ: Term.typ) (bc: t)
    : Term.t * Term.typ * t
    =
    let rec add term typ bc =
        match add_one_implicit term typ bc with
        | None ->
            term, typ, bc
        | Some (term, typ, bc) ->
            add term typ bc
    in
    add term typ bc




let unify (act: Term.typ) (bc: t): t option =
    let req = required_type bc in
    Option.map
        (fun gh -> {bc with gh})
        (Uni.unify act req true bc.gh)



let unify_plus (term: Term.t) (bc: t): (Term.t * t) option =
    let rec uni term tp bc =
        let tp = key_normal tp bc
        and req = required_type bc in
        match Uni.unify tp req true bc.gh with
        | Some gh ->
            Some (Gamma_holes.expand term gh, {bc with gh})
        | None ->
            Option.(add_one_implicit term tp bc
                    >>= fun (term, tp, bc) ->
                    uni term tp bc)
    in
    uni term (type_of_term term bc) bc




let set_term (term: Term.t) (bc: t): t =
    {bc with
        gh =
            Gamma_holes.fill_hole (index_of_level bc.sp bc) term bc.gh}



let make (gamma: Gamma.t): t =
    let cnt0 = Gamma.count gamma in
    {
        gh =
            Gamma_holes.(
                make gamma
                |> push_hole Term.(any_uni 2)
                |> push_hole Term.(Variable 0)
            );

        sp = cnt0 + 1;

        stack = [];

        entry = {
            cnt0;
        };

        entries = [];
   }




let final
    (bc: t)
    : (Term.t * Term.typ, int list * Term.typ * Gamma.t) result
    =
    let cnt0 = count_base bc
    and nlocs = count_locals bc in
    assert (bc.stack = []);
    assert (bc.entries = []);
    assert (bc.sp = cnt0 + 1);
    let term = top_term bc in
    let typ  = type_of_term term bc in
    match Term.down nlocs term, Term.down nlocs typ with
    | Some term, Some typ ->
        Ok (term, typ)
    | _ ->
        let gamma = Gamma_holes.context bc.gh in
        Error(
            unfilled_holes typ bc,
            typ,
            gamma
        )




let candidate
    (term: Term.t) (nargs: int) (bc: t)
    : (t, type_in_context * type_in_context) result
    =
    if 0 < nargs then
        let tp = type_of_term term bc in
        let term, tp, bc = add_implicits term tp bc in
        match unify tp bc with
        | None ->
            Error (required_type_in_context bc, type_in_context tp bc)
        | Some bc ->
            let bc = set_term term bc in
            let sp, stack = Stack.swap bc.sp bc.stack in
            Ok {bc with sp; stack}
    else
        match unify_plus term bc with
        | None ->
            Error (
                required_type_in_context bc,
                type_in_context (type_of_term term bc) bc
            )
        | Some (term, bc) ->
            Ok (set_term term bc)



let base_candidate
    (term: Term.t)
    (nargs: int)
    (bc: t)
    : t option
    =
    let term = Term.up (count_locals bc) term in
    match candidate term nargs bc with
    | Ok bc ->
        Some bc
    | Error _ ->
        None



let bound
    (ibound: int) (nargs: int) (bc: t)
    : (t, type_in_context * type_in_context) result
    =
    candidate (Gamma_holes.variable_of_bound ibound bc.gh) nargs bc





let next_formal_argument (name: string) (typed: bool) (bc: t): t =
    let cnt0 = count bc
    and tp = top_term bc
    in
    {bc with
        gh =
            Gamma_holes.(
                push_bound name typed tp bc.gh
                |> push_hole Term.(any_uni 1)
            );

        stack = bc.sp :: bc.stack;

        sp = cnt0 + 1
    }




module Product =
(* ... A: Any1, x: A, B: Any1, y: B, ... , RT: Any1 *)
struct
    let start (bc: t): t =
        let cnt0 = count bc
        in
        {
            gh = Gamma_holes.push_hole Term.(any_uni 1) bc.gh;

            stack = bc.sp :: bc.stack;

            sp = cnt0;

            entries = bc.entry :: bc.entries;

            entry = {cnt0}
        }

    let  check
        (nbounds: int)
        (bc: t)
        : (t, int) result
        =
        assert (0 < nbounds);
        let arg_tps, _ =
            let rec args nbounds stack tps =
                if nbounds = 0 then
                    tps, stack
                else
                    let sp, stack = Stack.split stack in
                    args (nbounds - 1) stack (term_at_level sp bc :: tps)
            in
            args nbounds bc.stack []
        in
        let rec find_incomplete i tps =
            if i = nbounds then
                None
            else
                let tp, tps = Stack.split tps in
                if Int_set.is_empty
                    (Gamma_holes.unfilled_holes bc.entry.cnt0 tp bc.gh)
                then
                    find_incomplete (i + 1) tps
                else
                    Some i
        in
        match find_incomplete 0 arg_tps with
        | Some i ->
            Error i
        | None ->
            Ok bc


    let end_
        (nargs: int)
        (nbounds: int)
        (bc: t)
        : (t, type_in_context * type_in_context) result
        =
        assert (0 < nbounds);
        let res = top_term bc
        and sp, stack = Stack.pop (nbounds + 1) bc.sp bc.stack
        in
        let tp = Gamma_holes.pi nbounds res bc.gh in
        let entry, entries = Stack.split bc.entries in
        candidate
            tp
            nargs
            {
                gh = Gamma_holes.remove_bounds nbounds bc.gh;
                sp;
                stack;
                entry;
                entries;
            }
end






module Typed =
(* [exp : tp] *)
struct
    let start (bc: t): t =
        (* Expect the type term. *)
        let cnt0 = count bc in
        {bc with
            gh = Gamma_holes.push_hole Term.(any_uni 1) bc.gh;

            stack = bc.sp :: bc.stack;

            sp = cnt0
        }

    let expression (bc: t): t =
        (* Expect the expression. *)
        let cnt0 = count bc in
        {bc with
            gh = Gamma_holes.push_hole (top_term bc) bc.gh;

            stack = bc.sp :: bc.stack;

            sp = cnt0;
        }

    let end_
        (nargs: int) (bc: t)
        : (t, type_in_context * type_in_context) result
        =
        let tp, stack = Stack.split bc.stack in
        let sp, stack = Stack.split stack in
        let tp = term_at_level tp bc
        and exp = top_term bc
        in
        let term = Term.Typed (exp, tp)
        and bc = {bc with sp; stack} in
        candidate term nargs bc
end






module Application =
(* [f a b c ...]

    f   a   b   c
    . g .
    .    h  .


For each argument we introduce 4 holes:

    A: Any(1)
    F: A -> Any(1)
    a: A
    f: all (x: A): F x
    ...

    stack = f a g b h c ...

*)
struct
    let start (nargs: int) (bc: t): t =
        let cnt0 = count bc in
        let rec shift cnt0 n gh sp stack =
            if n = 0 then
                gh, sp, stack
            else
                let open Gamma_holes in
                shift
                    (cnt0 + 4)
                    (n - 1)
                    (push_hole Term.(any_uni 1) gh  (* A: Any(1) *)
                    |> push_hole                    (* F: A -> Any(1) *)
                        Term.(Pi (Variable 0, any_uni 1, Pi_info.arrow))
                    |> push_hole Term.(Variable 1)  (* a: A *)
                    |> push_hole Term.(             (* f: all (x:A): F x *)
                        Pi (Variable 2,
                            Appl (
                                Variable 2,
                                Variable 0,
                                Application_info.Normal),
                            Pi_info.typed "x")))
                    (cnt0 + 3)
                    (cnt0 + 2 :: sp :: stack)
        in
        let gh, sp, stack = shift cnt0 nargs bc.gh bc.sp bc.stack in
        {bc with
            gh;
            stack;
            sp}


    let apply
        (n_remaining: int)
        (mode: Term.Application_info.t)
        (bc: t)
        : (t, type_in_context * type_in_context) result
        =
        match bc.stack with
        | f :: e :: stack ->
            let open Term in
            let term = Appl (
                    term_at_level f bc,
                    term_at_level bc.sp bc,
                    mode)
            in
            candidate
                term
                n_remaining
                {bc with
                    stack;
                    sp = e}
        | _ ->
            assert false (* Illegal call! *)
end (* Application *)





module Lambda =
(*
    \ (x: A) (y: B) ... : T := t

One placeholder for each optional argument type, one placeholder for the
optional result type and one placeholder for the inner expression.

    stack = t T ... B A ...
*)
struct
    let start (bc: t): t =
    (* Push a placeholder for the first argument type. I.e. expect the first
    argument type. *)
        let cnt0 = count bc in
        {bc with
            gh = Gamma_holes.push_hole Term.(any_uni 1) bc.gh;

            stack = bc.sp :: bc.stack;

            sp = cnt0
        }


    let inner (bc: t): t =
    (* Push a placeholder based on the result type (the previously analyzed ast),
    i.e. expect the inner term of the function abstraction. *)
        let cnt0 = count bc
        and tp = top_term bc
        in
        let open Gamma_holes in
        {bc with
            gh = push_hole tp bc.gh;

            stack = bc.sp :: bc.stack;

            sp = cnt0
        }


    let end_
        (nargs: int) (nbounds: int) (typed: bool) (bc: t)
        : (t, type_in_context * type_in_context) result
        =
        let exp = top_term bc in
        let sp, stack = Stack.split bc.stack in
        let exp =
            if typed then
                Term.Typed (exp, term_at_level sp bc)
            else
                exp
        in
        let lam = Gamma_holes.lambda nbounds exp bc.gh in
        let sp, stack = Stack.pop (nbounds + 1) sp stack
        in
        candidate
            lam
            nargs
            {bc with
                gh = Gamma_holes.remove_bounds nbounds bc.gh;
                sp;
                stack
            }
end





module Where =
struct
    let start (name: string) (bc: t): t =
        Application.start 1 bc
        |> Lambda.start
        |> next_formal_argument name false
        |> Lambda.inner



    let end_inner (bc: t): (t, type_in_context * type_in_context) result =
        Lambda.end_ 1 1 false bc


    let end_
        (nargs: int)
        (bc: t)
        : (t, type_in_context * type_in_context) result
        =
        match bc.stack with
        | f :: e :: stack ->
            let open Term in
            (
                match term_at_level f bc with
                | Lambda (tp, exp, info) ->
                    let term =
                        Where (
                            Lambda_info.name info,
                            tp,
                            exp,
                            term_at_level bc.sp bc
                        )
                    in
                    candidate term nargs {bc with stack; sp = e}
                | _ ->
                    assert false (* Illegal call! *)
            )
        | _ ->
            assert false (* Illegal call! *)
end