Type inference and program synthesis from simply typed lambda calculus type checking rules

stlc.prolog

?- set_prolog_flag(occurs_check,true).

lookup([(X,A)|_],X,A).
lookup([(Y,_)|T],X,A) :- \+ X = Y, lookup(T,X,A).

/* Rules from the STLC lecture */

pred(D,DD) :- D >= 0, DD is D - 1.

type(_,u,unit,D) :- pred(D,_).
type(G,app(M,N),B,D) :- pred(D,DD), type(G,M,A -> B,DD),type(G,N,A,DD).
type(G,lam(var(X),M),A -> B, D) :-
  pred(D,DD), type([(X,A)|G],M,B, DD).
type(G,fst(M),A,D) :-
  pred(D,DD), type(G,M,A * _,DD).
type(G,snd(M),B,D) :-
  pred(D,DD), type(G,M,_ * B,DD).
type(G,pair(M,N),A * B,D) :-
  pred(D,DD), type(G,M,A,DD), type(G,N,B,DD).
type(G,var(X),A,D) :-
  pred(D,_), lookup(G,X,A).

/* More rules for writing programs */
type(_,X,int,D) :-
  pred(D,_), integer(X).
type(_,D,int,D) :-
  pred(D,_).
type(_,true,bool,D) :-
  pred(D,_).
type(_,false,bool,D) :-
  pred(D,_).
type(G,A + B,int,D) :-
  pred(D,DD), type(G,A,int,DD), type(G,B,int,DD).
type(G,A < B,bool,D) :-
  pred(D,DD), type(G,A,int,DD), type(G,B,int,DD).
type(G,ite(A,B,C),T,D) :-
  pred(D,DD), type(G,A,bool,DD), type(G,B,T,DD), type(G,C,T,DD).

/* A generator for numbers */

gen(S,_,S).
gen(S,E,P) :- S < E, S2 is S+1, gen(S2,E,P).

/* Use iterative deepening */

synthesize(P,T) :-
  gen(0,10,D), type([],P,T,D).

/* Don't care about depth for inference. Use a large depth. */

infer(P,T) :-
  type([],P,T,1000).


/* Some queries to try */
/*
?- infer(1+2,T).
?- X = var(x), Y = var(y), infer(lam(X,lam(Y,ite(X < Y,X+Y,X-Y))),T).
?- X = var(x), infer(lam(X,fst(X)+1),T).
?- F = var(f), X = var(x), infer(lam(F,lam(X,app(F,X))),T).

?- X = var(x), infer(lam(X,app(X,X)),T).
?- infer(ite(true,0,false),T).

?- synthesize(P,int).
?- synthesize(P,(A*B)->A), var(A), var(B).
*/