type-inference.prolog

find(write, _, [], [], []).
find(_,    X, [X | After], [],           After).
find(Mode, X, [Y | Rest],  [Y | Before], After)
    :- find(Mode, X, Rest, Before, After).

access(Mode, Label, Pred, I, O) :-
    find(Mode, Label - X, I, B, A),
    !,
    call(Pred, X, Y),
    append([B, [Label - Y], A], O).

at(Label, Pred) --> access(write, Label, Pred).
look(Label, Pred) --> access(read, Label, Pred).

remove(Label, I, O) :-
    find(read, Label - _, I, B, A),
    !,
    append(B, A, O).

put(X, _, X).
get(X, X, X).

die(Msg) --> { writeln(Msg), false }.

dump(Stage) --> { writeln(Stage) }.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

kind(Ty, Kind) -->
    (   { Ty = app(F, X) }
    ->  !
    ,   kind(X, KX)
    ,   kind(F, KF)
    ,   ( { ShouldBe = (KX -> Kind) }
        , equal_kinds(KF, ShouldBe)

        ; die(cant_apply((F : KF) * (X : KX) \= Kind))
        )

    ;   { atom(Ty) }
    ->  at(kind, at(Ty, get(Kind)))
    ).

equal_kinds(K, K) --> {}.
equal_kinds(forall(_, Expr), K) -->
    { copy_term(Expr, Expr1) },
    equal_kinds(Expr1, K).
equal_kinds(K, forall(_, Expr)) -->
    { copy_term(Expr, Expr1) },
    equal_kinds(K, Expr1).

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

equal_types(L, R) -->
    (   { L = R }
    ->  !

    ;   { L = app(F, X), R = app(G, Y) }
    ->  !
    ,   kind(L, K)
    ,   kind(R, K)
    ,   equal_types(F, G)
    ,   equal_types(X, Y)

    ;   die(L \= R)
    ).

declare(Var, Ty) -->
    at(types, at(Var, get(poly(Ty)))).

undeclare(Var) -->
    at(types, remove(Var)).

declare_mono(Var, Ty) -->
    at(types, at(Var, get(mono(Ty)))).

lookup(Var, Ty) -->
    look(types, at(Var, get(GTy))),
    (   { GTy = poly(XTy) }
    ->  { copy_term(XTy, Ty) }

    ;   { GTy = mono(Ty) }
    ).

infer_check(lam(X, Body), app(app(->, XTy), BTy)) -->
    declare_mono(X, XTy),
    infer_check(Body, BTy),
    at(types, remove(X)).

infer_check(app(F, X), Ty) -->
    infer_check(F, FTy),
    equal_types(FTy, app(app(->, XTy), Ty)),
    infer_check(X, XTy).

infer_check(var(X), Ty) -->
    lookup(X, Ty).

infer_check(let(Ctx, Body), Ty) -->
    with_infer_check_bindings(Ctx, infer_check(Body, Ty)).

with_infer_check_bindings([], Pred) -->
    call(Pred).
with_infer_check_bindings([Name: Ty = Body | Rest], Pred) -->
    infer_check(Body, Ty),
    declare(Name, Ty),
    with_infer_check_bindings(Rest, Pred),
    undeclare(Name).

:-
    phrase(
        (   kind(int,  *)
        ,   kind(list, * -> *)

        ,   declare(head, app(app(->, app(list,A)), app(list,A)))

        ,   infer_check(let([twice : _ = lam(f, lam(x, app(var(f), app(var(f), var(x)))))], app(var(twice), var(head))), QTy)

        ,   declare(qhead, QTy)
        )
    , [], Ctx)
,   writeln(context: Ctx).