Generate simply-typed lambda calculus expressions of a given size

gen.pl

:- ensure_loaded(library(aggregate)).
:- ensure_loaded(library(lists)).

% Generate simply-typed lambda calculus terms of a given size.
%
% Based on Paul Tarau's paper A Hiking Trip Through the Orders of Magnitude:
% Deriving Efficient Generators for Closed Simply-Typed Lambda Terms and Normal
% Forms.
%
% The paper:
% https://arxiv.org/pdf/1608.03912.pdf
%
% An application of the algorithm shown here:
% https://github.com/pedrominicz/sugar

% Note that there are multiple definitions for the size of a lambda term. In
% particular, this definition is different from the one in Jue Wang's paper
% Generating Random Lambda Calculus Terms.
size(v(_), 0).
size(l(_, X), S) :- size(X, S0), S is S0 + 1.
size(a(X, Y), S) :- size(X, S0), size(Y, S1), S is S0 + S1 + 1.

n2s(0, z) :- !.
n2s(N, s(X)) :- N > 0, N0 is N - 1, n2s(N0, X).

down(s(X), X).

% Generate simply-typed lambda terms of a given size.
lambda(N, X) :- n2s(N, S), lambda(X, _, [], S, z).

lambda(v(X), A, Ctx) -->
  { member(X : A0, Ctx), unify_with_occurs_check(A0, A) }.
lambda(l(X : A, Y), A -> B, Ctx) -->
  down,
  lambda(Y, B, [X : A | Ctx]).
lambda(a(X, Y), B, Ctx) -->
  down,
  lambda(X, A -> B, Ctx),
  lambda(Y, A, Ctx).

% Count the number of simply-typed lambda terms of a given size.
count(N, Count) :-
  aggregate_all(count, lambda(N, _), Count).

% Print the counts of simply-typed lambda terms of at most a given size.
%
%   ?- counts(10).
%   1: 1
%   2: 2
%   3: 9
%   4: 40
%   5: 238
%   6: 1564
%   7: 11807
%   8: 98529
%   9: 904318
%   10: 9006364
%   true.
%
counts(Max) :-
  between(1, Max, N),
  count(N, Count),
  format('~d: ~d~n', [N, Count]),
  fail.
counts(_).

% Pretty print a term.
pretty(X) :-
  numbervars(X, 0, _),
  pretty(X, Xs, []),
  maplist(write, Xs),
  nl.

pretty(v('$VAR'(I))) --> [x, I].
pretty(l('$VAR'(I) : A, X)) -->
  ['(λ ', x, I, ' : '], pretty_type(A), [', '], pretty(X), [')'].
pretty(a(X, Y)) --> ['('], pretty(X), [' '], pretty(Y), [')'].

pretty_type(A -> B) -->
  ['('], pretty_type(A), !, [' → '], pretty_type(B), [')'].
pretty_type(A) --> [A].

% Print all simply-typed lambda terms of a given size.
show(N) :-
  lambda(N, X),
  pretty(X),
  fail.
show(_).