"complete and easy" (mostly), in ELPI

t.elpi

kind ty type.

type tunit  ty.
type tfun   ty -> ty -> ty.
type tall   (ty -> ty) -> ty.

kind tm type.

type unit   tm.
type app    tm -> tm -> tm.
type lam    (tm -> tm) -> tm.
type annlam ty -> (tm -> tm) -> tm.
type annot  tm -> ty -> tm.

% necessary? maybe.
type mono ty -> o.
mono (tfun X Y) :- mono X, mono Y.

type sub ty -> ty -> o.
sub X Y :- var X, !, instL X Y.
sub X Y :- var Y, !, instR X Y.
% now, neither argument is a variable
sub tunit tunit.
sub (tfun Ax Bx) (tfun Ay By) :- !, sub Ay Ax, sub Bx By.
sub X (tall Y) :- !, pi a\ (mono a => sub X (Y a)).
sub (tall X) Y :- !, sub (X _) Y.
sub X X.

% the left argument is a variable, gotta instantiate it to a monotype
type instL ty -> ty -> o.
instL X Y :- var Y, !, X = Y.
instL (tfun Ax Bx) (tfun Ay By) :- !, instR Ay Ax, instL Bx By.
instL X (tall Y) :- !, pi a\ (mono a => instL X (Y a)).
instL Y Y :- mono Y.

% the right argument is a variable, gotta instantiate it to a monotype
type instR ty -> ty -> o.
instR X Y :- var X, !, X = Y.
instR (tfun Ax Bx) (tfun Ay By) :- !, instL Ay Ax, instR Bx By.
instR (tall X) Y :- !, instR (X Fresh_) Y.
instR X X :- mono X.

% the mutually recursive type inference functions
type check  tm -> ty -> o.
type syn    tm -> ty -> o.
type apply  ty -> tm -> ty -> o.

check E T :- var T, !, syn E A, sub A T.
check E (tall Ty) :- !, pi a\ (mono a => check E (Ty a)).
check (lam F) (tfun A B) :- !, pi x\ (syn x A => check (F x) B).
check (annlam A F) (tfun A_ B) :- !, sub A_ A, pi x\ (syn x A => check (F x) B).
check E T :- syn E A, sub A T.

syn (annot E T) T :- !, check E T.
syn (app F X) B :- !, syn F A, apply A X B.
syn unit tunit :- !.
syn (lam F) (tfun A B) :- !, pi x\ (syn x A => syn (F x) B).
syn (annlam A F) (tfun A B) :- !, pi x\ (syn x A => syn (F x) B).

apply T X B :- var T, !, T = tfun A B, syn X A.
apply (tall T) X B :- !, apply (T Fresh_) X B.
apply (tfun A B) X B :- check X A.

% Simple example with pairs and units
type tunit   ty.
type tpair  ty -> ty -> ty.
mono tunit.
mono (tpair X Y) :- mono X, mono Y.

type pair tm.
syn pair (tall a\ tall b\ tfun a (tfun b (tpair a b))).

type example tm -> o.
% example ((λ f : (∀a. a -> a). pair (f unit) (f f)) (λ x. x))
example (app (annlam (tall a\ tfun a a) f\ app (app pair (app f unit)) (app f f)) (lam x\ x)).

% Output of elpi t.elpi <<<'example X, syn X T.'
% Success:
%   T = tpair tunit (tfun X0 X0)

% ST example: runST has a rank-2 type
type st     ty -> ty -> ty.
type stref  ty -> ty -> ty.
mono (st S T) :- mono S, mono T.
mono (stref S T) :- mono S, mono T.

type runST tm.
syn runST (tall a\ tfun (tall s\ st s a) a).

type newSTRef tm.
syn newSTRef (tall a\ tall s\ tfun a (st s (stref s a))).

% GHC gives the error 'type variable s would escape its scope'
type badexample tm -> o.
badexample (app runST (app newSTRef unit)).

% Output of elpi t.elpi <<<'badexample X, syn X T.'
% Failure