\section{Proofs}\label{app:proofs}

\begin{proposition}
  The language of compressed terms is strongly normalizing.
\end{proposition}

\begin{proof}
  Consider evaluating $\hat{t}$ in applicative order (leftmost innermost).
  %
  In this case, any left-hand size of a $\beta$-reduction has no inner $\beta$-redexes.
  %
  Since in our language a $\lambda$-abstraction can only be on the left-hand side of a $\beta$-redex,
  it means that expression being reduced has no inner $\lambda$-abstractions.
  %
  For that reason, the number of $\lambda$-abstractions in an expression strictly decreases with every $\beta$-step
  (the sole one in the reduced redex disappears, and all other $\lambda$-abstractions are outside of the reduced redex, and hence unchanged).
\end{proof}

\begin{lemma}
Given a $\kappa$-step $\hat{t}_1 \rwsteps{\kappa(p')} \hat{t}_2$
  if a pattern $p$ has $N$ matches in $\hat{t}_2$, it also has at least $N$ matches in $\hat{t}_1$.
  \label{lem:pat-match}
\end{lemma}
\begin{proof}
  There are four cases for what $p$ can match in $\hat{t}_2$:
  \begin{itemize}
  \item a subterm that does not overlap with the newly introduced $\beta$-redex: in this case, the match is unaffected by the $\kappa$-step;
  \item a subterm inside an actual argument of the $\beta$-redex: this actual argument becomes a sub-term of $\hat{t}_1$;
  \item a sub-term inside the body of the $\beta$-redex: $\hat{t}_1$ has a sub-term that is more specific than the body, and hence still matches $p$;
  \item a sub-term that includes the entire $\beta$-redex inside: since $p$ is first order, it cannot mention the redex, so the redex must be contained entirely inside the substitution;
  hence again the match is unaffected by the $\kappa$-step.
  \end{itemize}
\end{proof}

\begin{theorem}[Soundness and Completeness of Pattern-Based Library Learning]
  For any term $t \in \termsof{\Sigma}$ and compressed term $\hat{t} \in \absof{\Sigma}{\mathcal{X}}$:
  \begin{description}[font=\itshape]
    \item[(Soundness)] If $t$ compresses into $\hat{t}$, then $\hat{t}$ evaluates to $t$: $\forall \mathcal{P} . t \rewrites{\kappa(\mathcal{P})} \hat{t} \Longrightarrow \hat{t} \rewrites{\beta} t$.
    \item[(Completeness)] If $\hat{t}$ is a solution to the (global) library learning problem,
    then $t$ compresses into $\hat{t}$ using only patterns that have a match in $t$:
    $\hat{t} \in \argmin_{\hat{t}' \rewrites{\beta} t} \mathsf{size}(\hat{t}') \Longrightarrow t \rewrites{\kappa(\mathcal{P})} \hat{t}$,
    where $\mathcal{P} = \{p \in \patsof{\Sigma}{\mathcal{X}} \mid t' \in \mathsf{subterms}(t) , t' \sqsubseteq p\}$.
  \end{description}
\label{thm:pat-lib-learning-app}
\end{theorem}

\begin{proof}
The soundness is trivial because inverting any $\kappa$-rewrite gives a valid $\beta$-reduction.

The other direction (completeness) is more involved.
First let us prove that $\hat{t}$ can be obtain by any compression, regardless of whether the patterns occur in $t$.
This is non-trivial
because inverting a $\beta$-reduction $(\lambda \many{X} \ \mathbf{\shortrightarrow}\  \hat{t}_1)\ \many{\hat{t}_2}  \rwsteps{\beta}  \hat{t}'$
does not always correspond to a well-form $\kappa$-step,
for two reasons:
(a) $\hat{t}_1$ itself contains $\lambda$-abstractions (and hence does not correspond to a pattern);
(b) $\many{X} \neq \mathsf{vars}(\hat{t}_1)$.
To overcome point (a), note that in an applicative evaluation,
both $\hat{t}_1$ and $\hat{t}_2$ contain no $\beta$-redexes (and hence no $\lambda$-abstractions).
Hence  $\hat{t}_1$ is a valid pattern.

For point (b), there are two cases: either (1) $\mathsf{vars}(\hat{t}_1) \subsetneq \many{X}$
or (2) $\mathsf{vars}(\hat{t}_1) \supsetneq \many{X}$.
Consider case (1).
In this case, the compressed term $\hat{t}$ has $\lambda$-bindings that are not used in their bodies.
Such a term cannot possibly be minimal.
Removing unused bindings makes each individual $\beta$-redex smaller and cannot remove sharing,
so it always makes the overall term smaller as well.

Now consider case (2).
This means that $\lambda \many{X} \ \mathbf{\shortrightarrow}\  \hat{t}_1$ has free variables,
which must be defined in outer $\lambda$-abstractions.
However, using such a $\lambda$-abstraction contradicts the assumption that we are only interested in global library learning.

Finally, let us prove that all patterns used in the compression have a match in $t$.
Consider a compression $t = \hat{t}_0 \rwsteps{\kappa(p_1)} \hat{t}_1 \dots \rwsteps{\kappa(p_n)}  \hat{t}_n = \hat{t}$.
Clearly, each pattern $p_i$ has a match in the (compressed) term $\hat{t}_{i-1}$, by definition of the $\kappa$-rewrite.
Hence, it also occurs in $t$, by induction using Lemma~\ref{lem:pat-match}.
\end{proof}

\newpage
\section{Additional Experiments}\label{app:eval}

\begin{figure}[h]
  \centering
  \begin{subfigure}{0.48\linewidth}
    \includegraphics[width=\linewidth]{figures/scatter/list-rounds.pdf}
    \caption{List domain}
  \end{subfigure}
  \hfill
  \begin{subfigure}{0.48\linewidth}
    \includegraphics[width=\linewidth]{figures/scatter/physics-rounds.pdf}
    \caption{Physics domain}
  \end{subfigure}

  \begin{subfigure}{0.48\linewidth}
    \includegraphics[width=\linewidth]{figures/scatter/text-rounds.pdf}
    \caption{Text domain}
  \end{subfigure}
  \hfill
  \begin{subfigure}{0.48\linewidth}
    \includegraphics[width=\linewidth]{figures/scatter/logo-rounds.pdf}
    \caption{Logo domain}
  \end{subfigure}

  \begin{subfigure}{0.48\linewidth}
    \includegraphics[width=\linewidth]{figures/scatter/towers-rounds.pdf}
    \caption{Towers domain}
  \end{subfigure}
  \caption{
    \tname{babble}'s default configuration runs for 20 rounds (r=20)
    to learn additional library functions.
    This plot shows the data from \autoref{fig:dc-domains},
     but with additional \tname{babble} configuration running fewer rounds.
    With fewer rounds, \tname{babble} runs faster but compresses worse.
  }\label{fig:dc-domains-rounds}
\end{figure}
\section{Library Learning over Terms}\label{sec:au}

In this section
  we formalize the problem of
  library learning over a corpus of terms and
  motivate our first core contribution:
  generating candidate abstractions via anti-unification.
\autoref{sec:egraphs} generalizes this formalism to library learning over an e-graph.
For simplicity of exposition,
our formalization of library learning is \emph{first order},
that is, the initial corpus does not itself contain any $\lambda$-abstractions,
and all the learned abstractions are first-order functions
(the \tname{babble} implementation does not have this limitation).

\subsection{Preliminaries}

\mypara{Terms}
A \emph{signature} $\Sigma$ is a set of constructors,
each associated with an arity.
$\termsof{\Sigma}$ denotes the set of \emph{terms} over $\Sigma$,
defined as the smallest set containing all $s(t_1,\dots,t_k)$
where $s \in \Sigma$, $k = \mathsf{arity}(s)$, and $t_1,\dots,t_k\in\termsof{\Sigma}$.
We abbreviate nullary terms of the form $s()$ as $s$.
The \emph{size} of a term $\mathsf{size}(t)$ is defined in the usual way
(as the number of constructors in the term).
We use $\mathsf{subterms}(t)$ to denote the set of all subterms of $t$ (including $t$ itself).
We assume that $\Sigma$ contains a dedicated \emph{variadic} tuple constructor,
written $\langle t_1, \ldots, t_n \rangle$,
which we use to represent a corpus of programs as a single term
($\langle \rangle$ does not contribute to the size of a term).

\mypara{Patterns}
If $\mathcal{X}$ is a denumerable set of variables,
$\patsof{\Sigma}{\mathcal{X}}$ is a set of \emph{patterns},
\emph{i.e.}\xspace terms that can contain variables from $\mathcal{X}$.
A pattern is \emph{linear} if each of its variables occurs only once:
$\forall X \in \mathsf{vars}(p) . \mathsf{occurs}(X, p) = 1$.
A \emph{substitution} $\sigma = [\subst{X_1}{p_1}, \ldots, \subst{X_n}{p_n}]$
is a mapping from variables to patterns.
We write $\sapp{\sigma}{p}$ to denote the application of $\sigma$ to pattern $p$,
which is defined in the standard way.
We define the size of a substitution $\mathsf{size}(\sigma)$ as the total size of its right-hand sides.

A pattern $p$ is more general than (or \emph{matches}) $p'$,
written $p' \sqsubseteq p$,
if there exists $\sigma$ such that $p' = \sapp{\sigma}{p}$;
we will denote such a $\sigma$ by $\match{p'}{p}$.
For example $X + 1 \sqsubseteq X + Y$ with $\match{X + 1}{X + Y} = [\subst{Y}{1}]$.
The relation $\sqsubseteq$ is a partial order on patterns,
and induces an equivalence relation $p_1 \sim p_2 \triangleq p_1 \sqsubseteq p_2 \wedge p_2 \sqsubseteq p_1$
(equivalence up to variable renaming).
In the following, we always distinguish patterns only up to this equivalence relation.

The \emph{join} of two patterns $p_1 \sqcup p_2$---also called their \emph{anti-unifier}---%
is the least general pattern that matches both $p_1$ and $p_2$;
the join is unique up to $\sim$.
Note that $\langle \patsof{\Sigma}{\mathcal{X}}, \sqsubseteq, \sqcup, \top = X \rangle$ is a join semi-lattice
(part of Plotkin's \emph{subsumption lattice}~\cite{plotkin1970lattice}).
Consequently, the join can be generalized to an arbitrary set of patterns.

A \emph{context} $\mathcal{C}$ is a pattern with a single occurrence of a distinguished variable $\circ$.
We write $\mathcal{C}[p]$---$p$ in context $\mathcal{C}$---as a syntactic sugar for $\sapp{[\subst{\circ}{p}]}{\mathcal{C}}$.
A \emph{rewrite rule} $R$ is a pair of patterns, written $p_1 \Rightarrow p_2$.
Applying a rewrite rule $R$ to a term or pattern $p$, written $\sapp{R}{p}$ is defined in the standard way:
$\sapp{R}{p} = \sapp{(\match{p}{p_1})}{p_2}$ if $p \sqsubseteq p_1$ and undefined otherwise.
A pattern $p$ can be \emph{re-written in one step} into $q$ using a rule $R$,
written $p \rwsteps{R} q$,
if there exists a context $\mathcal{C}$ such that $p = \mathcal{C}[p']$ and $q = \mathcal{C}[\sapp{R}{p'}]$.
The reflexive-transitive closure of this relation is the \emph{rewrite relation} $\rewrites{\mathcal{R}}$,
where $\mathcal{R}$ is a set of rewrite rules.

\mypara{Compressed Terms}
Compressed terms $\absof{\Sigma}{\mathcal{X}}$ are of the form:
$$
\hat{t} ::= X   \mid   s(\hat{t}_1,\ldots,\hat{t}_k)     \mid (\lambda X_1 \ldots X_n \ \mathbf{\shortrightarrow}\  \hat{t})\ \hat{t}_1\ \ldots\ \hat{t}_n
$$
In other words,
compressed terms may contain applications of a $\lambda$-abstraction with zero or more binders
to the same number of arguments.
Note that this is a first-order language in the sense that all abstractions are fully applied.

Importantly, we define $\mathsf{size}(\hat{t})$ in such a way
that multiple occurrences of a $\lambda$-abstraction are \emph{only counted once}.
For simplicity of accounting, the size of an application $(\lambda \many{X} \ \mathbf{\shortrightarrow}\  \hat{t}_1)\ \many{\hat{t}_2}$
is defined as $\mathsf{size}(\hat{t}_1) + \sum{\many{\mathsf{size}(\hat{t}_2)}}$,
that is, abstraction nodes themselves do not add to the size.%
\footnote{Hereafter, $\many{a}$ stands for a sequence of elements of syntactic class $a$ and $\epsilon$ denotes the empty sequence.}

\emph{Beta-reduction} on compressed terms, denoted $\hat{t}_1 \rwsteps{\beta} \hat{t}_2$, is defined in the usual way:
\[
  (\lambda \many{X} \ \mathbf{\shortrightarrow}\  \hat{t}_1)\ \many{\hat{t}_2}  \rwsteps{\beta} \sapp{[\many{X \mapsto \hat{t}_2}]}{\hat{t}_1}
\]
where substitution on compressed terms is the standard capture-avoiding substitution for $\lambda$-calculus.
Note that because our language is first-order and has no built-in recursion,
it is strongly normalizing
(the proof of this statement, as well as other proofs omitted from this section, can be found in the supplementary material).
Hence, the order of $\beta$-reductions is irrelevant,
and without loss of generality we can define \emph{evaluation} on compressed terms, $\hat{t}_1 \rewrites{\beta} \hat{t}_2$,
to follow applicative order, \emph{i.e.}\xspace innermost $\beta$-redexes are reduced first.
As a result, any application reduced during evaluation
has the form $(\lambda \many{X} \ \mathbf{\shortrightarrow}\  p_1)\ \many{p_2}$,
that is, neither the body $p_1$ nor the actual arguments $\many{p_2}$ contain any redexes
(and hence any $\lambda$-abstractions).
This simplifies several aspects of our formalization;
for example, there is no need for $\alpha$-renaming,
since with no binders in $p_1$, no variable capture can occur.

\mypara{Problem Statement}
We can now formalize
  the \emph{library learning problem} as follows:
  given a term $t \in \termsof{\Sigma}$,
  the goal is to find the smallest compressed term $\hat{t} \in \absof{\Sigma}{\mathcal{X}}$
  that evaluates to $t$ (\emph{i.e.}\xspace $\hat{t} \rewrites{\beta} t$).
The reason such $\hat{t}$ may be smaller than $t$,
  is that it may contain multiple occurrences of the same $\lambda$-abstraction
  (applied to different arguments),
  whose size is only counted once.
An example is shown
  in \autoref{fig:lib-learning-example}.

Although in full generality the solution might include nested lambdas with free variables
(defined in the outer lambdas),
in the rest of the paper we restrict our attention to \emph{global library learning},
where all lambdas are closed terms.
This is motivated by the purpose of library learning
to discover reusable abstraction for a given problem domain.
The solution in \autoref{fig:lib-learning-example} already has this form.

\begin{figure}
  \begin{minipage}{.5\textwidth}
    \begin{align*}
      \langle &f(g(a) + g(a)) + (g(1) + h(2))\\
      ,\ &f(g(b) + g(b)) + (g(3) + h(4))\\
      ,\ &g(5) + h(6) \rangle
    \end{align*}
  \end{minipage}%
  \begin{minipage}{.5\textwidth}
    \begin{align*}
      &\langle \hat{f}_2\ g(a)\ 1\ 2, \hat{f}_2\ g(b)\ 2\ 3, \hat{f}_1\ 5\ 6\rangle\quad \text{where}\\
      &\hat{f}_1 = \lambda Y\ Z \ \mathbf{\shortrightarrow}\  g(Y) + h(Z)\\
      &\hat{f}_2 = \lambda X\ Y\ Z \ \mathbf{\shortrightarrow}\  f(X + X) + \hat{f}_1\ Y\ Z
    \end{align*}
  \end{minipage}
  \caption{Library learning.
           Left: initial term (size 29).
           Right: an optimal solution with two abstractions, one of which uses the other (size 26).
           A solution with $\hat{f}_2 = \lambda X\ Y\ Z \ \mathbf{\shortrightarrow}\  f(g(X) + g(X)) + \hat{f}_1\ Y\ Z$ also has size 26.}\label{fig:lib-learning-example}
\end{figure}

\subsection{Pattern-Based Library Learning}

At a high-level, our approach to library learning
is to use \emph{patterns} that occur in the original corpus
as candidate bodies for $\lambda$-abstractions in the compressed corpus.
Looking at the example in \autoref{fig:lib-learning-example},
it is not immediately obvious that using just the patterns from the original corpus is sufficient,
since the body of $f_2$ contains an application of $f_1$.
Perhaps surprisingly, this is not an issue:
the key idea is that we can compress $t$ into $\hat{t}$ by inverting the evaluation $\hat{t} \rewrites{\beta} t$,
and because the evaluation order is applicative,
the rewritten sub-term at every step will not contain any $\beta$-redexes.

\mypara{Compression}
More formally, given a pattern $p$,
let us define its \emph{compression rule} (or $\kappa$-rule for short)
as the rewrite rule
\[
  \kappa(p) \; \triangleq \; p \Rightarrow (\lambda \many{X} \ \mathbf{\shortrightarrow}\  p)\ \many{X} \quad \text{where}\quad \many{X} = \mathsf{vars}(p)
\]
In other words, $\kappa(p)$ will replace any term
matching $p$ with an application of a function whose body is \emph{exactly} $p$.
For example, if $p = g(Y) + h(Z)$,
its $\kappa$-rule is $g(Y) + h(Z) \Rightarrow (\lambda Y\ Z \ \mathbf{\shortrightarrow}\  g(Y) + h(Z))\ Y\ Z$.
Note that on the right-hand side of this rule
only the \emph{free} occurrences of $Y$ and $Z$ will be substituted during rewriting;
the bound $Y$ and $Z$ will be left unchanged, following the usual semantics of substitution for $\lambda$-calculus.
For example, this rule can rewrite the third term in \autoref{fig:lib-learning-example} (left) as follows:
\[
  g(5) + h(6) \;\rwsteps{\kappa(g(Y) + h(Z))}\; (\lambda Y\ Z \ \mathbf{\shortrightarrow}\  g(Y) + h(Z))\ 5\ 6
\]
A sequence of $\kappa$-rewrites $t \rwsteps{\kappa(p_1)} \dots \rwsteps{\kappa(p_n)} \hat{t}$,
where all $p_i \in \mathcal{P}$,
is called a \emph{compression} of $t$ into $\hat{t}$ using patterns $\mathcal{P}$
and written $t \rewrites{\kappa(\mathcal{P})} \hat{t}$.
We can now show that the library learning problem is equivalent to finding the smallest compression of $t$
using only patterns that occur in $t$.

\begin{theorem}[Soundness and Completeness of Pattern-Based Library Learning]
  For any term $t \in \termsof{\Sigma}$ and compressed term $\hat{t} \in \absof{\Sigma}{\mathcal{X}}$:
  \begin{description}[font=\itshape]
    \item[(Soundness)] If $t$ compresses into $\hat{t}$, then $\hat{t}$ evaluates to $t$: $\forall \mathcal{P} . t \rewrites{\kappa(\mathcal{P})} \hat{t} \Longrightarrow \hat{t} \rewrites{\beta} t$.
    \item[(Completeness)] If $\hat{t}$ is a solution to the (global) library learning problem,
    then $t$ compresses into $\hat{t}$ using only patterns that have a match in $t$:
    $\hat{t} \in \argmin_{\hat{t}' \rewrites{\beta} t} \mathsf{size}(\hat{t}') \Longrightarrow t \rewrites{\kappa(\mathcal{P})} \hat{t}$,
    where $\mathcal{P} = \{p \in \patsof{\Sigma}{\mathcal{X}} \mid t' \in \mathsf{subterms}(t) , t' \sqsubseteq p\}$.
  \end{description}
\label{thm:pat-lib-learning}
\end{theorem}
The proof can be found in the supplementary material.

\mypara{Example}
Consider once again the library learning problem in \autoref{fig:lib-learning-example}.
Here the set of patterns used to compress the original corpus into the solution on the right is:
\[
  p_1 = g(Y) + h(Z)\quad\quad\quad
  p_2 = f(X + X) + g(Y) + h(Z)
\]
Rewriting the first term of the corpus proceeds in two steps (the redexes of $\kappa$-steps are highlighted):
\begin{multline*}
  \textcolor{darkblue}{f(g(a) + g(a)) + (g(1) + h(2))}    \;\rwsteps{\kappa(p_2)}\\
  (\lambda X\ Y\ Z \ \mathbf{\shortrightarrow}\  f(X + X) + \textcolor{darkblue}{g(Y) + h(Z)})\ g(a)\ g(a)\ (g(1) + h(2))\;\rwsteps{\kappa(p_1)}\\
  (\lambda X\ Y\ Z\ \ \mathbf{\shortrightarrow}\  f(X + X) + (\lambda Y\ Z \ \mathbf{\shortrightarrow}\  g(Y) + h(Z))\ Y\ Z)\ g(a)\ g(a)\ (g(1) + h(2))
\end{multline*}
In other words, we first rewrite the entire term using $p_2$,
and then rewrite inside the body of the introduced abstraction using $p_1$
(note that this order of compression is the inverse of the applicative evaluation order).
The second term of the corpus compresses analogously;
the third term compresses in a single step using $p_1$.
Although this is not obvious from the rewrite sequence above,
the resulting compressed corpus is indeed smaller than the original
thanks to sharing of both $\lambda$-abstractions,
as illustrated on the right of \autoref{fig:lib-learning-example}.

\mypara{Library Learning as Term Rewriting}
Theorem~\ref{thm:pat-lib-learning} reduces library learning to a \emph{term rewriting} problem.
Namely, given a term $t$ and a finite set of rewrite rules $\mathcal{R} = \{\kappa(p) \mid  t' \in \mathsf{subterms}(t) , t' \sqsubseteq p\}$,
our goal is to find a minimal-size term $\hat{t}$ such that $t \rewrites{\mathcal{R}} \hat{t}$,
which is a standard formulation in term rewriting.
Unfortunately, this particular problem is notoriously difficult because
(a) the rule set $\mathcal{R}$ is very large for any non-trivial term $t$, and
(b) our $\mathsf{size}$ function is non-local (it takes sharing into account)
In the rest of this section we discuss how we can prune the rule set $\mathcal{R}$
to reduce it to a tractable size.
\autoref{sec:beam} discusses how we tackle the remaining term rewriting problem
using the \emph{equality saturation} technique~\cite{tate2009equality,willsey2021egg}.

\subsection{Pruning Candidate Patterns}
\label{sec:pruning}

In this section, we discuss which patterns can be discarded from consideration
when constructing the set of $\kappa$-rules $\mathcal{R}$ for the term rewriting problem.

\mypara{Cost of a Pattern}
Consider a compression $t \rewrites{\kappa(\mathcal{P})} \hat{t}$
where each pattern $p \in \mathcal{P}$ is used some number $n$ times,
with substitutions $\sigma^p_1, \ldots, \sigma^p_n$.
We can break down the total amount of compression into contributions of individual patterns:
\[
  \mathsf{size}(\hat{t}) - \mathsf{size}(t) = \sum_{p \in \mathcal{P}} \mathsf{cost}(p, \{\sigma^p_1, \dots \sigma^p_n\})
\]
The cost of a pattern $p$, in turn, consists of three components.
The cost of \emph{introducing} the abstraction is the size of its body, \emph{i.e.}\xspace $\mathsf{size}(p)$.
The cost of using an abstraction---%
$\textsf{use}(p, \sigma)$ \eqref{eq:use}---%
 includes the application itself and the size of the arguments.
The cost saved by using an abstraction---%
$\textsf{save}(p, \sigma)$ \eqref{eq:save}---%
 is just the cost of the term matched by $p$
 (\emph{i.e.}\xspace the redex of the corresponding $\kappa$-step).
\begin{align}

  \textsf{use}(p, \sigma) &= 1 + \sum_{X \in \mathsf{vars}(p)} \mathsf{size}(\sigma(X)) \label{eq:use} \\
  \textsf{save}(p, \sigma) &= \mathsf{size}(\sapp{\sigma}{p}) =
  \mathsf{size}(p) + \sum_{X \in \mathsf{vars}(p)} \mathsf{occurs}(X, p) \cdot (\mathsf{size}(\sigma(X)) - 1)\label{eq:save}
\end{align}
The total cost of $p$ is
 the cost of introducing the abstraction paid a single time,
 plus the cost of each use, minus what you save for each application:
\begin{equation}
  \textsf{cost}(p, \{\sigma_1,\ldots,\sigma_n\})
  = \sz(p) + \sum_{\sigma_i} (\textsf{use}(p, \sigma_i) - \textsf{save}(p, \sigma_i))
\end{equation}
When $p$ is linear (all $\mathsf{occurs}(X, p) = 1$),
 the cost depends only on $n$ but not on the substitutions $\sigma_i$:
\begin{align}
  \textsf{cost}(p, \{\sigma_1,\ldots,\sigma_n\})
  &= \sz(p) + \sum_{\sigma_i} \left(
    1 - \mathsf{size}(p) + |\mathsf{vars}(p)|
  \right) \\
  &= \sz(p) + n \cdot \left(
    1 - \mathsf{size}(p) + |\mathsf{vars}(p)|
  \right)
\end{align}
We can show that a pattern $p$ with a \emph{non-negative} cost can be safely discarded,
that is: there exists another compression using only $\mathcal{P} \setminus \{p\}$,
whose result is at least as small.

\mypara{Trivial Patterns}
Based on this analysis, any linear pattern $p$
 with $\mathsf{skeleton}(p) \leq 1$ can be discarded,
 where $\mathsf{skeleton}(p) = \mathsf{size}(p) - |\mathsf{vars}(p)|$ is the size of $p$'s ``skeleton'', \emph{i.e.}\xspace it's body without the variables.
Intuitively, the skeleton of $p$ is simply too small to pay for introducing an application.
In this case, $\mathsf{cost}(p,\_) > 0$ \emph{independently} of how many times it is used.
We refer to such patterns as \emph{trivial}.
Examples of trivial patterns are $X$ and $X + Y$.

\mypara{Patterns with a Single Match}
We can show that patterns with only a single match in the corpus can also be discarded.
If $p$ has a single match in $t$,
then it can appear \emph{at most once} in any compression of $t$.
If $p$ is linear, $\mathsf{cost}(p, \_) = 1 + |\mathsf{vars}(p)|$,
which is always positive, so $p$ can be discarded.
But what about non-linear patterns, where even a single $\kappa$-step can decrease size thanks to variable reuse?
It turns out that any non-linear pattern with a single match
can always be replaced by a nullary pattern (with no variables) that is more optimal.

Without loss of generality, assume that $p$ has a single variable $X$ that occurs $m > 1$ times,
and let its sole $\kappa$-step be $\sapp{\sigma}{p} \rwsteps{\kappa(p)}  (\lambda X \ \mathbf{\shortrightarrow}\  p)\  \sapp{\sigma}{X}$.
The size of the right-hand side is $\mathsf{size}(p) + 1 + \mathsf{size}(\sapp{\sigma}{X})$,
or, rewritten in terms of $p$'s skeleton: $1 + (\mathsf{skeleton}(p) + m) + \mathsf{size}(\sapp{\sigma}{X})$.
Instead, we can rewrite the same redex $\sapp{\sigma}{p} \rewrites{\{R\}} p'$
using $m$ applications of the rule $R \;\triangleq\; \sapp{\sigma}{X} \Rightarrow (\lambda \epsilon \ \mathbf{\shortrightarrow}\  \sapp{\sigma}{X})\ \epsilon$.
This is a $\kappa$-rule for a \emph{nullary} pattern $\sapp{\sigma}{X}$ with no variables
(hence the corresponding $\lambda$-abstraction has zero binders).
The size of $p'$ obtained in this way is $\mathsf{size}(\sapp{\sigma}{X}) + m + \mathsf{skeleton}(p)$
(where the former is the size of the shared $\lambda$-abstraction, $m$ is the number of applications,
and $\mathsf{skeleton}(p)$ is the size of the term around the applications).
As you can see, this term is one smaller than the one we get by applying $p$.
Intuitively, this result says that instead of using a non-linear pattern that occurs only once,
it is better to perform common sub-expression elimination.

\mypara{Parameterization Lattice}
Eliminating from consideration all patterns with fewer than two matches in the corpus
suggests an algorithm for generating a complete set $\mathcal{P}$ of candidate patterns:
  (1) start from the set $\plau{t} = \{t_1 \sqcup t_2 \mid t_i \in \mathsf{subterms}(t)\}$
      of all pairwise joins of subterms of the input program,
  (2) explore all elements of the subsumption semi-lattice above those patterns,
      by gradually replacing sub-patterns with variables,
      until we hit trivial patterns at the top of the lattice.
We will refer to this semi-lattice above $\plau{t}$ as the \emph{parametrization lattice} of $t$,
denoted $\pl{t}$.
A fragment of $\pl{t}$ for $t = \langle f(a+b), f(a+c), f(b+c)\rangle$ is shown in \autoref{fig:lattice} (left).

\begin{figure}
  \begin{minipage}{.5\textwidth}
    \centering
    \includegraphics[width=.8\textwidth]{figures/lattice-1.pdf}
  \end{minipage}%
  \begin{minipage}{.5\textwidth}
    \centering
    \includegraphics[width=.8\textwidth]{figures/lattice-2.pdf}
  \end{minipage}
  \caption{Left: fragment of the parametrization lattice for the term $t = \langle f(a+b), f(a+c), f(b+c)\rangle$;
           only filled black circles correspond to candidate patterns:
           hollow circles match fewer than two terms and gray circles are trivial.
           Right: an example where is it insufficient to consider pairwise joins
           to obtain an optimal pattern.}
  \label{fig:lattice}
\end{figure}

\mypara{Approximation}
In practice, computing the set $\plau{t}$ is feasible:
although there are quadratically many pairs of subterms,
most of them do not have a common constructor at the root, and hence their join is trivially $X$.
An example is the join of $t$ with any of its subterms in \autoref{fig:lattice} (left).
Unfortunately, generalizing the patterns from $\plau{t}$
according to the parameterization lattice (\autoref{fig:lattice}) is expensive.
For this reason, \tname{babble} adopts an approximation
and simply uses $\plau{t}$ as the set of candidates.

This approximation makes our pattern generation theoretically incomplete.
Consider a pattern $p \in \pl{t}\setminus \plau{t}$;
there are two reasons why we might need $p$ in the optimal compression of $t$:
\begin{enumerate}
  \item there is \emph{no} $p' \in \plau{t}$ with the same set of matches as $p$, or
  \item there is such a $p'$ but it has \emph{few enough matches}
  that its larger size does not pay off.

\end{enumerate}

The first kind of incompleteness occurs when $p$ matches a set of subterms $t_1, \ldots, t_n$ ($n > 2$),
whose join is distinct from all their pairwise joins
(otherwise some $t_i \sqcup t_j \in \plau{t}$ would also match all $t_1, \ldots, t_n$).
An example is shown in \autoref{fig:lattice} (right),
where the three subterms in question are $f(a+a)$, $f(a+c)$, and $f(c+c)$.
In this case, an optimal compression might use the pattern $f(X+Y)$ to rewrite all three subterms,
but our approximation would only include the patterns $f(a+X)$, $f(X+X)$, and $f(X+c)$,
each of which can rewrite only two of the three subterms.

The second kind of incompleteness occurs when there exists $p' \in \plau{t}$
that has the same set of matches as $p$, despite being strictly more specific,
and yet using $p'$ instead of $p$ still doe not pay off.
The understand when this happens, consider the difference in costs between $p$ and $p'$,
assuming that they are both used to rewrite the same $n$ subterms
(\emph{i.e.}\xspace their \textsf{save} cost is the same):
\begin{align*}
  \mathsf{cost}(p', \many{\sigma'}) - \mathsf{cost}(p, \many{\sigma}) &= \mathsf{size}(p') - \mathsf{size}(p) + \sum_i (\mathsf{use}(p', \sigma'_i) - \mathsf{use}(p, \sigma_i))\\
  &= \mathsf{size}(p') - \mathsf{size}(p) + \sum_i (\mathsf{size}(\sigma'_i) - \mathsf{size}(\sigma_i))
\end{align*}
Because $p'$ is strictly more specific than $p$, we know that $\mathsf{size}(p') \geq \mathsf{size}(p)$,
but all its substitutions $\sigma'_i$ must be strictly smaller than $\sigma_i$.
Hence, with enough uses, $p'$ is bound to become more compressive than $p$;
when there are just a few uses, however, $p$ can be more optimal.
For example, consider the corpus $\langle \mathcal{C}[f(1,2,3)],  \mathcal{C}[f(4,5,6)]\rangle$,
where $\mathcal{C}$ is some sufficiently large context.
Here, a more general pattern $p = \mathcal{C}[X]$ is more optimal than the more specific $p' = \mathcal{C}[f(X,Y,Z)]$,
because $\mathsf{size}(p') - \mathsf{size}(p) = 3$, each use of $p'$ is only one node cheaper,
and there are only two uses.

Despite the lack of theoretical completeness guarantees,
we argue that restricting candidate patterns to $\plau{t}$ is a reasonable trade-off.
Note, that the counter-examples above are quite contrived,
and they no longer apply once the corpus contains sufficiently many and sufficiently diverse instances of a pattern
(for example, adding $f(b,b)$ to the first corpus would make $f(X,Y)$ appear in $\plau{t}$,
and adding just one more occurrence of $p'$ into the second corpus would make it as optimal as $p$).
Our empirical evaluation confirms that this approximation works well in practice.

\section{Candidate Selection via E-Graph Extraction}\label{sec:beam}

After generating candidate abstractions,
  the \T{LLMT} algorithm invokes \T{select_library} to pick the subset of candidate patterns
  that can best be used to compress the input corpus.
This section describes our approach to selecting the optimal library dubbed
  \emph{targeted common subexpression elimination}.

\mypara{Library Selection as E-Graph Extraction}
Recall that candidate selection starts with an e-graph $\mathcal{G}'$,
  which represents all the ways of compressing the initial corpus and its equivalent terms
  using the candidate patterns $\mathcal{P}$.
We will refer to a subset $\mathcal{L} \subset \mathcal{P}$ as a \emph{library}.
The optimal size of an e-class $c$ compressed with $\mathcal{L}$
  can be computed as the sum of the sizes of
  (1) the smallest term $t\in \denot{c}$
      using only the library functions in $\mathcal{L}$ and
      where $\lambda$-abstractions
      do not count toward the size of $t$, and
  (2) the smallest version of each $p \in \mathcal{L}$.

Note that
  the cost of defining an abstraction
  in $\mathcal{L}$ is only counted once, and
  that abstractions in $\mathcal{L}$ can be used
  to compress other abstractions in $\mathcal{L}$.
Our goal is to find $\mathcal{L}$ such that the root e-class compressed with $\mathcal{L}$
  has the smallest size.

Given a \emph{particular} library,
  we can find the size of the smallest term
  via a relatively straightforward top-down traversal of the e-graph.
Hence, a na\"ive approach to library selection
  would be to enumerate all subsets of $\mathcal{P}$
  and pick the one that produces the smallest term at the root.
Unfortunately, this approach becomes intractable as the size of $\mathcal{P}$ grows.

\begin{figure}
  \textbf{E-node cost set} $\cs_N(n)$
  $$
  \begin{array}{rl}
    \cs_N(s()) &= \{ (\emptyset, 1) \} \\
    \cs_N(s(\many{a_i})) &= \{ (\mathcal{L}, u + 1) \mid (\mathcal{L}, u) \in \mathsf{cross}(\many{a_i}) \} \\

    \cs_N((\lambda \many{X} \ \mathbf{\shortrightarrow}\  a)\ b) &= \mathsf{addlib}(a, \cs(a), \cs(b)) \\
    \end{array}
  $$

  \textbf{E-class cost set} $\cs(a)$
  $$
  \begin{array}{rl}
    \cs(\{ \many{n_j} \}) &= \ensuremath{\textcolor{darkred}{\mathsf{prune}}}(\ensuremath{\textcolor{darkblue}{\mathsf{reduce}}}(\bigcup \many{\cs_N(n_j)})) \\
  \end{array}
  $$

  \textbf{Auxiliary definitions} $\mathsf{cross}$, $\mathsf{addlib}$, $\ensuremath{\textcolor{darkblue}{\mathsf{reduce}}}$, $\ensuremath{\textcolor{darkred}{\mathsf{prune}}}$
  $$
  \begin{array}{rl}
    \mathsf{cross}(z_1, z_2) &= \ensuremath{\textcolor{darkred}{\mathsf{prune}}}(\ensuremath{\textcolor{darkblue}{\mathsf{reduce}}}(\{ (\mathcal{L}_1 \cup \mathcal{L}_2, u_1 + u_2) \mid (\mathcal{L}_1, u_1) \in z_1, (\mathcal{L}_2, u_2) \in z_2 \})) \\
    \mathsf{addlib}(a, z_1, z_2) &= \ensuremath{\textcolor{darkred}{\mathsf{prune}}}(\ensuremath{\textcolor{darkblue}{\mathsf{reduce}}}(\{ (\mathcal{L}_1 \cup \mathcal{L}_2 \cup \{ a \}, u_2) \mid (\mathcal{L}_1, u_1) \in z_1, (\mathcal{L}_2, u_2) \in z_2 \})) \\
    \ensuremath{\textcolor{darkblue}{\mathsf{reduce}}}(z) &= \{ (\mathcal{L}_1, u_1) \in z \mid \forall (\mathcal{L}_2, u_2) \in z \setminus (\mathcal{L}_1, u_1).\ \mathcal{L}_1 \subset \mathcal{L}_2 \lor u_1 < u_2 \} \\
    \ensuremath{\textcolor{darkred}{\mathsf{prune}}}(N, K, z) &=  \ensuremath{\mathsf{top\_k}}(\{(\mathcal{L}, u) \in z \mid |\mathcal{L}| \leq N \}, K)\\
  \end{array}
  $$
\caption{Cost set propagation, defined as two mutually
  recursive functions. Blue text corresponds to the partial order
  reduction optimization and red text corresponds to the beam approximation.
  $\ensuremath{\mathsf{top\_k}}(S, K)$ is a helper function that returns top $K$
  elements from the sorted set $S$.}
  \label{fig:costset}
\end{figure}

\mypara{Exploiting Partial Shared Structure}
Instead, \tname{babble} selects the optimal library using a bottom-up dynamic programming algorithm.
To this end, it associates each e-node and e-class with a \emph{cost set},
which is a set of pairs $(\mathcal{L}, u)$
  where $\mathcal{L}$ is a library and $u$ is the \emph{use cost} of this library,
  \emph{i.e.}\xspace the size of the smallest term represented by the e-node / e-class if it is allowed to use $\mathcal{L}$
  (excluding the size of $\mathcal{L}$ itself).
Cost sets are propagated up the e-graph using the rules shown in \autoref{fig:costset}.
The base case is a nullary e-node $s()$,
  which cannot use any library functions and whose size is always $1$.
For an e-node that has children,
\tname{babble} computes the $\mathsf{cross}$ product over the cost sets of all its child e-classes\xspace.
Finally, for an application e-node $(\lambda \many{X} \ \mathbf{\shortrightarrow}\  a)\ b$,
  the cost set \emph{must include} the library function $a$
  (in addition to some combination of libraries from the cost sets of $a$ and $b$);
  note that the use cost of the abstraction node only includes the use cost of $b$,
  since abstraction bodies are excluded from the use cost.

To compute the cost set of an e-class,
  \tname{babble} takes the union of the cost sets of all its e-nodes.
However, doing this naively would result in the size of the cost set growing exponentially.
To mitigate this, we define a \emph{partial order reduction} ($\ensuremath{\textcolor{darkblue}{\mathsf{reduce}}}$),
which only prunes provably sub-optimal cost sets.
Given two pairs $(\mathcal{L}_1, u_1)$ and $(\mathcal{L}_2, u_2)$ in the cost set of an e-class,
if $\mathcal{L}_1 \subset \mathcal{L}_2$ and $u_1 \leq u_2$,
then $\mathcal{L}_2$ is subsumed by $\mathcal{L}_1$ and can be removed from the cost set,
intuitively because $\mathcal{L}_1$ can compress this e-class even better and with fewer library functions.
In practice this optimization prunes libraries with redundant abstraction,
where two different abstractions can be used to compress the same subterms.

\mypara{Beam Approximation}
Even with the partial order reduction,
  calculating the cost set for every e-node and e-class
  can blow up exponentially.
To mitigate this,
  \tname{babble} provides the option to limit
  both the \emph{size of each library} inside a cost set
  and the \emph{size of the cost set} stored for each e-class.
This results in a \emph{beam-search} style algorithm,
where the cost set of each e-class is $\ensuremath{\textcolor{darkred}{\mathsf{prune}}}$ed, as shown in \autoref{fig:costset}.
This pruning operation first filters out libraries that have more than $N$ patterns,
then ranks the rest by the total cost (\emph{i.e.}\xspace the sum of use cost and the size of the library),
and finally returns the top $K$ libraries from that set.

\section{Candidate Selection vis E-Graph Extraction}\label{sec:beam}

In this section we formalize \emph{beam extraction}, our approach
to selecting an optimal subset of abstractions
from a wider set of candidate abstractions in the presence of an e-graph.

\subsection{Library Selection as E-Graph Extraction}

Having generated potential candidates for learned abstractions
through e-graph anti-unification, the problem of LLMT now becomes
\emph{candidate selection}: which subset of learned abstractions
should we use to optimally rewrite our target term?
In the presence of e-graphs, this answer to this question is
non-trivial. As mentioned in \autoref{sec:egraphs}, not only can
$\denot{\mathcal{G}}$ be potentially infinite, making enumerating terms in
$\mathcal{G}$ impractical, but the choice of which subset of abstractions
is optimal depends on the concrete term itself.

To tackle this problem, we reduce the problem of library selection to
that of extracting a single, optimal concrete term from an e-graph.
Similarly to how we define a correspondence between patterns and
compression rules, which rewrite an expression into an application
of an abstraction,
we can also define a correspondence between patterns and
\emph{library rules}, which rewrite an expression into a combined
abstraction definition and application.
For a given \T{let}-binding $\T{let}\ P = \lambda \many{X} \to p$,
let its library rule $\mathsf{lib}(P, \many{X}, p)$ be the rewrite rule
$p \Rightarrow \T{let}\ P = \lambda \many{X} \to p\ \T{in}\ P\ \many{X}$.
As an example, $\mathsf{lib}(P, X, X + 1)$ corresponds to the rewrite rule
$X + 1\Rightarrow \T{let}\ P = \lambda X \to X + 1\ \T{in}\ P\ X$.
By generating a rewrite $\mathsf{lib}(P_i, \mathsf{vars}(p_i), p_i)$ for every pattern
$p_i$ generated during e-graph anti-unification of an e-graph $\mathcal{G}$,
and applying these rewrites via equality saturation to $\mathcal{G}$, we
obtain an e-graph $\mathcal{G}'$ in which every possible application to a new
learned abstraction is included in the e-graph, always accompanied by an associated
abstraction definition at every application.
Because every application of a learned abstraction in any subterm of
any $t \in \denot{\mathcal{G}'}$ must be accompanied by a corresponding
\T{let}-binding, the subset of abstractions used in $t$ can be determined
solely by its \T{let}-bindings. Thus, finding the optimal set of libraries
to use for $\mathcal{G}'$ can be found by extracting the smallest term in
$\mathcal{G}'$, where ``smallest'' is given by a cost function $\cs$
which only counts the definitions of \T{let}-bindings once.
The set of abstractions which are \T{let}-bound in this term will be
the optimal subset of libraries to use for the programs represented by
$\mathcal{G}'$.

\subsection{Optimal Extraction with Partial Shared Structure}

Having reduced library subset selection to the problem of extracting
a term from an e-graph in the presence of shared structure, the problem
now becomes defining a cost function $\cs$ which can simultaneously
account for common \T{let}-defined sub-expressions while also allowing
for efficient optimal term extraction.
Traditional e-graph term extraction won't work in this context, since
this cost function cannot be \TODO{not monotonic but another word
goes here}, and greedy term extraction
inherently relies on having a \TODO{same} cost function.
For instance, consider the following e-graph with eight e-classes,
corresponding to the term $f(g(a, 1, a), g(b, 1, a), g(c, 1, c))$
after the library rewrite $\mathsf{lib}(l, x, g(x, 1, x))$ has been applied:
\begin{align*}
c_0&\colon \{ a \},\ c_1\colon \{ b \},\ c_2\colon \{ c \},\ c_3\colon \{ 1 \} \\
c_4&\colon \{ g(c_0, c_3, c_0), \T{let}\ l = \lambda x \to g(x, 1, x)\ \T{in}\ l(c_0) \} \\
c_5&\colon \{ g(c_1, c_3, c_1), \T{let}\ l = \lambda x \to g(x, 1, x)\ \T{in}\ l(c_1) \} \\
c_6&\colon \{ g(c_2, c_3, c_2), \T{let}\ l = \lambda x \to g(x, 1, x)\ \T{in}\ l(c_2) \} \\
c_7&\colon \{ f(c_4, c_5, c_6) \} \\
\end{align*}
A purely greedy extraction approach based on AST size would dictate that
$g(a, 1, a)$ should be the selected subterm at $c_4$, since it has the smallest
AST size of all the expressions in this e-class.
However, when used in the context of $c_7$, if the definition of $f$
is shared, then $\T{let}\ l = \lambda x \to g(x, 1, x)\ \T{in}\ l(a)$
should be selected, since the cost of the extra \T{let}-binding
can be amortized by its later use in $c_7$.
Indeed, in general, greedy extraction would \emph{never} extract a
use of an abstraction, since the cost of defining and applying
an abstraction always exceeds the cost of \TODO{the expression
itself} at the point of definition;
this cost only becomes \TODO{worth it} if the abstraction is
used at differet parts of the e-graph, which \TODO{a greedy
extraction algorithm wouldn't be able to surmise}.

Because our common sub-expression extraction problem is
inherently limited---only definitions in a \T{let}-binding are
considered shared---we can define a cost function which keeps track
of \T{let}-defined symbols and counts them once, while counting all
other subterms without regard for common sub-expressions.
By separating common sub-expression cost from \TODO{regular node?}
cost, we can \TODO{find the true cost of a term by finding
the sum of libs and expr cost.}

\begin{figure}
  \textbf{E-node cost set} $\cs_N(n)$
  $$
  \begin{array}{rl}
    \cs_N(s()) &= \{ (\emptyset, 1) \} \\
    \cs_N(s(a)) &= \{ (l, c + 1) \mid (l, c) \in \cs(a) \} \\
    \cs_N(s(\many{a_i})) &= \mathsf{cross}(\many{\cs(a_i)}) \\
    \cs_N(\T{let}\ s = \lambda \many{X} \to a\ \T{in}\ b) &= \{ (l, c + 1) \mid (l, c) \in \mathsf{clib}(s, \cs(a), \cs(b)) \} \\
    \end{array}
  $$

  \textbf{E-class cost set} $\cs(a)$
  $$
  \begin{array}{rl}
    \cs(\{ \many{n_j} \}) &= \bigcup \many{\cs_N(n_j)} \\
  \end{array}
  $$

  \textbf{Argument cost set combination} $\mathsf{cross}$
  $$
  \begin{array}{rl}
    \mathsf{cross}(z_1, z_2) &= \{ (l_1 \cup l_2, c_1 + c_2) \mid (l_1, c_1) \in z_1, (l_2, c_2) \in z_2 \} \\
  \end{array}
  $$

  \textbf{Abstraction cost set combination} $\mathsf{clib}$
  $$
  \begin{array}{rl}
    \mathsf{clib}(s, z_1, z_2) &= \{ (l_1 \cup l_2 \cup \{ (s, c_1) \}, c_2) \mid (l_1, c_1) \in z_1, (l_2, c_2) \in z_2 \} \\
  \end{array}
  $$
\caption{\TODO{caption for costset def}}\label{fig:costset}
\end{figure}

\autoref{fig:costset} describes the cost function calculation for
e-nodes and e-classes.
Every e-node and e-class has a corresponding \emph{cost set},
which represents a set of possible library selections for all terms
represented by the e-node or e-class:
a choice of a set of abstractions, their associated cost, and the
size of the rewritten term using these abstractions.
These cost sets are propagated bottom-up through the e-graph.

To justify our cost function definition, we revisit our previous
example, \TODO{in which greedy extraction would not have worked}.
For the e-class $c_0$, there is only one possible selection
of abstractions and one possible rewritten term cost for all the
terms in each e-class: the empty set of abstractions with term cost $1$.
This cost set, $\{ (\emptyset, 1) \}$, is the same as that of $c_1$,
$c_2$, and $c_3$.
To find the cost set for $c_4$, we must first find the cost set
calculations for each of its e-nodes.
Finding the cost set of the expression $g(c_0, c_3, c_0)$ requires
us to combine the cost sets of the argument e-classes $c_0$
and $c_3$;
in our formalization, this is performed by the $\mathsf{cross}$
operation.
The possible library subsets which can be used to rewrite $g(c_0, c_3, c_0)$
consists of any subset which can be used for $c_0$ combined
with any subset which can be used for $c_3$; the corresponding
term cost in each case is simply the term costs of $c_0$ and $c_3$ added
together (since this outer expression includes the subterms for $c_0$
and $c_3$) and incremented, to account for the $+$ AST node.
In this case, since the cost sets of $c_0$ and $c_3$ only contain one
element each, the cost set of $c_0 + c_3$ consists of one element:
$\{ (\emptyset \cup \emptyset, 1 + 1 + 1)\} = \{ (\emptyset, 3) \}$.
This corresponds to the intuition that the term $a + 1$ can be
expressed with the empty set of libraries, in which case it has term
size $3$.
For an e-node with three or more arguments, the $\mathsf{cross}$
operation is performed as a fold over the arguments of the e-node.

Finding the cost set for the other e-node in $c_4$,
$\T{let}\ f = \lambda x \to x + 1\ \T{in}\ f(c_0)$,
requires first finding the cost set for the application, $f(c_0)$.
The cost set of a unary e-node is equal to the cost set of
its argument, where all term costs have been incremented,
so this application would have a cost set of
$\{ (\emptyset, 2) \}$.
To find the cost set of this expression as a whole, we must
add the abstraction $f$ and its associated cost into every
possibility represented by the cost set of the \T{let}-binding's
body.
One complication with this is that the definition of the
abstraction itself is an e-class which can include applications
to other learned abstractions, and thus has its own cost set.
For the sake of simplicity, we assume the body of $f$ contains
one possible term, $\lambda x \to x + 1$, which would have
a cost set of $\{ (\emptyset, 3) \}$.
In this case, we see that the cost set of this expression should
be $\{ (\{ (f, 3) \}, 2) \}$, representing the fact that
this e-node can be expressed with term size 2 given an abstraction
$f$ of term size 3.
In general, this is handled by the $\mathsf{clib}$ operation,
which performs a similar function as $\mathsf{cross}$: for every
possible definition of an abstraction (and thus every possible
abstraction subset and term size it can have), it must be
combined with every possible abstraction and term size of the
\T{let}-binding's body.
In this case, ``combined'' means that all abstraction definitions,
including the \T{let}-bound abstraction itself and any other
abstractions it might use, are \TODO{unioned} together,
while the term sizes of the body are kept constant.
With the cost sets of these two e-nodes calculated, we can now
find the cost set for the e-class $c_4$ as a whole.
The cost set of an e-class is defined as the union of the cost sets
of its e-nodes;
the possible selections of library subsets in an e-class is equal to
all of the possible selections of library subsets of any of its
constituent e-nodes.
In this case, this means that
$\cs(c_4) = \{ (\emptyset, 3), (\{ (f, 3) \}, 2) \}\}$.
In other words, a term $t \in \denot{c_4}$ can have AST
size 3 if rewritten with no abstractions, or term size 2
if rewritten with the abstraction $f$ of size 3..
The cost sets of $c_5$ and $c_6$ are the same as $c_4$, and
the cost set for $c_7$ can be found in a similar fashion,
resulting in a final cost set:

\begin{align*}
\cs(c_7) = \{ &(\emptyset, 10), (\{ (f, 3) \}, 9), (\{ (f, 3) \}, 9), (\{ (f, 3) \}, 9), \\
&(\{ (f, 3) \}, 8), (\{ (f, 3) \}, 8), (\{ (f, 3) \}, 8), (\{ (f, 3) \}, 7) \}
\end{align*}

Library subset extraction is performed by examining the root
e-class and finding the library subset selection which results
in the lowest \emph{full cost}: the sum of the costs of the
shared expressions, the abstraction definitions, combined with
the term cost.
In this case, \TODO{that lowest is given by this.
Thus, the optimal library is this. Woo!}

\mypara{Partial Order Reduction}
The size of the cost set of an e-node is exponential with respect
to the number of arguments it has.
\TODO{This is not very performant. But we can prune a bunch of redundant
things. Consider the cost set from this last example, in which we define
a lib but never actually use it. we can remove this possibility completely!}

\mypara{From Optimal to Beam Extraction}
Even in the presence of partial order reduction, calculating this cost
set for every e-node and e-class in an e-graph representing a large
program can still be prohibitively expensive.
In practice, \tname{babble} provides multiple parameters which limit the number
of candidate libraries at each node (the \emph{beam size}),
as well as the size of each library.
\section{Conclusion and Future Work}

We presented library learning modulo theory (LLMT),
 a technique for learning abstractions from a corpus of programs
 modulo a user-provided equational theory.
We implemented LLMT in \tname{babble}.
Our evaluation showed that \tname{babble}
 achieves better compression orders of magnitude faster
 than the state of the art.
On a larger benchmark suite of \tname{2d cad} programs,
 \tname{babble} learns sensible functions that
 compress a dataset that was---until now---too large for library learning techniques.

LLMT and \tname{babble} present many avenues for future work.
First,
 our evaluation showed that equational theories are important
 for achieving high compression,
 but these must be provided by domain experts.
Recent work in automated theory synthesis like
 Ruler~\cite{ruler} or \tname{TheSy}~\cite{thesy}
 could aid the user in this task.
Second,
 LLMT uses e-graph anti-unification to generate
 promising abstraction candidates,
 but this approach is incomplete and misses some patterns
 that could achieve better compression.
An exciting direction for future work is to combine LLMT
with more efficient top-down search from \tname{Stitch}~\cite{stitch}.
This is challenging because \tname{Stitch} crucially relies on the ability
to quickly compute an upper bound on the compression of a given pattern
by summing up the local compression at each of its matches in the corpus.
This upper bound does not straightforwardly extend to e-graphs
because in an e-graph different matches of a pattern may come from different syntactic variants of the corpus,
and one needs to trade-off the compression from abstractions
against the size difference between different syntactic variants.

\section{Library Learning modulo Equational Theory}\label{sec:egraphs}

\begin{figure}
  \begin{minipage}{.5\textwidth}
    \textbf{Syntax}
    $$
    \begin{array}{rll}
      \text{e-class ids} & a,b &\in \mathcal{I} \\
      \text{e-nodes}     & n \ensuremath{\mathrel{\mathord{:}\mathord{:=}}} s(a_1, \ldots, a_k)  &\in N\\
      \text{e-classes}   & c \ensuremath{\mathrel{\mathord{:}\mathord{:=}}} \{n_1, \ldots, n_m\} &\in C
    \end{array}
    $$
  \end{minipage}%
  \begin{minipage}{.5\textwidth}
    \textbf{Denotation} $\denot{\cdot}\colon N\to 2^{\termsof{\Sigma}}$, $\denot{\cdot}\colon C\to 2^{\termsof{\Sigma}}$
    $$
    \begin{array}{rl}
      \\
      \denot{s(a_1, \ldots, a_k)} &= \{s(t_1, \ldots, t_k) \mid t_i \in \denot{M(a_i)}\}\\
      \denot{\{n_1, \ldots, n_m\}} &= \bigcup_{i \in [1,m]} \denot{n_i}
    \end{array}
    $$
  \end{minipage}
  \caption{Syntax, metavariables, and denotation for the components of an e-graph.
           Here $s \in \Sigma$ and $k = \mathsf{arity}(s)$.}\label{fig:syntax}
\end{figure}

\mypara{E-Graphs}
Let $\mathcal{I}$ be a denumerable set of \emph{e-class ids}.
An \emph{e-graph} $\mathcal{G}$ is a triple $\langle C, M, r\rangle$,
where $C$ is a set of \emph{e-classes},
$M\colon \mathcal{I} \to C$ is an \emph{e-class map}.
and $r\in \mathcal{I}$ is the root class id.
An \emph{e-class} $c \in C$ is a set of \emph{e-nodes} $n \in N$,
and an e-node is a constructor applied to e-class ids.
The syntax of e-classes and e-nodes is summarized in \autoref{fig:syntax} (left).
An e-graph has to satisfy the \emph{congruence invariant},
which states that the e-graph has no two identical e-nodes
(or alternatively, that all e-classes are disjoint).%
\footnote{In a real e-graph implementation, the definitions of e-graphs and the congruence invariant
are more involved,
because efficient merging of e-classes requires introducing a non-trivial equivalence relation over e-class ids;
these details are irrelevant for our purposes.
Also, other formalizations of e-graphs do not feature a distinguished root e-class.}

The \emph{denotation} of an \emph{e-graph}---the set of terms it represents---%
is the denotation of its root e-class $\denot{M(r)}$,
where the denotation of e-classes and e-nodes is defined mutually-recursively
in \autoref{fig:syntax} (right).
Note that the denotation can be infinite if the e-graph has cycles.
An e-graph induces an equivalence relation $\semeq^{\mathcal{G}}$,
where $t_1 \semeq^{\mathcal{G}} t_2$ iff there exists an e-class $c \in C$ such that $t_1 \in \denot{c} \wedge t_2\in \denot{c}$.

E-graphs provide means to \emph{extract} the cheapest term from an e-class according to some cost function:
$\mathsf{extract}_{\mathsf{cost}}(a) = \argmin_{t\in \denot{M(a)}} \mathsf{cost}(t)$.
If $\mathsf{cost}$ is \emph{local},
 meaning that the cost of a term can be computed
 from the costs of its immediate children,
extraction can be done efficiently by a greedy algorithm,
which recursively extracts the best term from each e-class.

\mypara{E-Matching}
E-matching is a generalization of pattern matching to e-graphs,
where matching an e-class $c$ against a pattern $p$ yields a set of \emph{e-class substitutions}
$\theta\colon \mathcal{X}\to \mathcal{I}$
such that $\sapp{\theta}{p}$ is a ``subgraph'' of $c$.
To formalize this notion, we introduce \emph{partial terms} $\pi \in \patsof{\Sigma}{\mathcal{I}}$,
which are terms whose leaves can be e-class ids
(or, alternatively, patterns with e-class ids for variables).
The containment relation $\contains{\pi}{a}$ for some e-class id $a$ is defined as follows:
$$
\contains{a}{a} \quad\quad\quad
\contains{s(\pi_1, \ldots, \pi_k)}{a}\ \text{iff}\ s(a_1, \ldots, a_k) \in M(a) \wedge \contains{\pi_i}{a_i}
$$
With this definition, a pattern $p$ matches an e-class $a$, $a \sqsubseteq p$, if there exists an e-class substitution $\theta$,
such that $\contains{\sapp{\theta}{p}}{a}$.
We denote the set of such substitutions as $\matches{a}{p}$.

\mypara{Rewriting and Equality Saturation}
Equality saturation (EqSat)~\cite{tate2009equality,willsey2021egg}
takes as input a term $t$ and a set of equations that induce an equivalence relation $\equiv$,
and produces an e-graph $\mathcal{G}$ such that $\denot{\mathcal{G}} = \{t' \mid t \equiv t'\}$
and $t \equiv t'$ iff $t \semeq^{\mathcal{G}} t'$.
The core idea of EqSat is to convert equations into rewrite rules
and apply them to the e-graph in a non-destructive way:
so that the original term and the rewritten terms are both represented in the same e-class.
Applying a rewrite rule $p_1 \Rightarrow p_2$ to an e-class $a$ works as follows:
for each $\theta \in \matches{a}{p_1}$,
we obtain the rewritten partial term $\pi' = \sapp{\theta}{p_2}$
and then add this partial term to the same e-class $a$,
restoring the congruence invariant
(\emph{i.e.}\xspace merging e-classes that now have identical e-nodes).

\subsection{Top-Level Algorithm}

\input{llmt_algorithm}

We can formalize the problem of \emph{library learning modulo equational theory} (LLMT) as follows:
given a term $t$ and a set of equations that induce an equivalence relation $\equiv$,
the goal is to find a compressed term $\hat{t} \in \absof{\Sigma}{\mathcal{X}}$,
such that $\hat{t} \rewrites{\beta} t' \equiv t$ (for some $t'$),
and $\hat{t}$ has a minimal size.

Our top-level algorithm \T{LLMT} is depicted in \autoref{fig:algo-llmt}.
This algorithm takes as input the original corpus $t$
and the equational theory, represented as a set of require rules $\mathcal{R}_\equiv$
(another input to the algorithm is the maximum size of the library;
this parameter is introduced for the sake of efficiency, as we explain in \autoref{sec:beam}).
As the first step (lines 4--5),
\T{LLMT} applies EqSat to obtain an e-graph $\mathcal{G}$
that represents all terms $t'$ such that $t' \equiv t$.
This reduces the LLMT problem to library learning over an e-graph:
\emph{i.e.}\xspace the goal is to find a minimal-size compressed term $\hat{t}$,
such that $\hat{t} \rewrites{\beta} t'$ for some $t' \in \denot{\mathcal{G}}$.
Similarly to \autoref{sec:au},
we take a pattern-based approach to this problem,
that is, we select a set $\mathcal{P}$ of \emph{candidate patterns}
and then perform compression rewrites using these patterns.

The rest of the algorithm is split into two phases:
\emph{candidate generation} and \emph{candidate selection}.
Candidate generation (lines 8--10)
first computes the set of candidate patterns $\mathcal{P}$ using the anti-unification mechanism extended to e-graphs.
Then it creates a compression rule ($\kappa$-rule, see \autoref{sec:au}) from each candidate pattern,
and once again applies EqSat to obtain a new e-graph $\mathcal{G}'$.
This new e-graph represents all possible ways to compress the terms from $\mathcal{G}$
using patterns in $\mathcal{P}$.
Finally, the candidate selection phase (lines 13--15)
selects the optimal subset of compression rules $\mathcal{R}'$,
constructs an e-graph $\mathcal{G}''$ that represents all possible compressions using only the selected compression rules,
and finally extracts the smallest compressed term from this e-graph.

The rest of this section focuses on the candidate generation via e-graph anti-unification
(line 8).
The candidate selection functions \T{select\_library} and \T{extract} are discussed in \autoref{sec:beam}.

\subsection{Candidate Generation via E-Graph Anti-Unification}

The goal of candidate generation is to find a set of patterns
that are useful for compression.
Following the discussion in \autoref{sec:au},
we restrict our attention to patterns $\mathsf{AU}(t_1, t_2)$,
where $t_1$ and $t_2$ are subterms of some $t \in \denot{\mathcal{G}}$.
The na\"ive approach is to enumerate all such $t$,
and for each one, perform anti-unification on all pairs of subterms;
this is suboptimal at best,
and impossible at worst (when $\denot{\mathcal{G}}$ is infinite).
Hence in this section we show how to compute a finite set of candidate patterns
directly on the e-graph,
without discarding any optimal patterns.

\mypara{E-Class Anti-Unification}
Let us first consider anti-unification of two e-classes, $\mathsf{AU}(a,b)$,
which takes as input e-class ids $a$ and $b$
and returns a set of patterns.
We define $\mathsf{AU}(a,b) = \ecau{\epsilon}{a}{b}$,
where $\ecau{\Gamma}{a}{b}$ is an auxiliary function
that additionally takes into account a \emph{context} $\Gamma$.
A context is a list of pairs of e-classes that have been visited while computing the AU,
and is required to prevent infinite recursion in case of cycles in the e-graph.

\begin{figure}
  \textbf{E-node Anti-Unification} $\ecau{\Gamma}{n_1}{n_2}$
  $$
  \begin{array}{rl}
    \ecau{\Gamma, (a,b)}{s(a_1,\ldots,a_k)}{s(b_1,\ldots,b_k)} &= \{s(p_1,\ldots,p_k) \mid p_i \in \ecau{\Gamma}{a_i}{b_i}\}\\
    \ecau{\Gamma, (a,b)}{s_1(\ldots)}{s_2(\ldots)} &= \{X_{a,b}\} \quad\text{if}\ s_1 \neq s_2\\
    \end{array}
  $$
  \textbf{E-class Anti-Unification} $\ecau{\Gamma}{a}{b}$
  $$
  \begin{array}{rl}
    \ecau{\Gamma}{a}{b} &= \emptyset \quad\text{if}\ (a,b) \in \Gamma\\

    \ecau{\Gamma}{a}{b} &= \textcolor{darkblue}{\mathsf{dominant}} \left(
       \bigcup_{n_a \in M(a), n_b \in M(b)} \ecau{\Gamma, (a, b)}{n_a}{n_b}
    \right)
  \end{array}
  $$
\caption{E-class anti-unification defined as two mutually-recursive functions of e-nodes and e-class ids.}\label{fig:ecau}
\end{figure}

\autoref{fig:ecau} defines this operation using two mutually recursive functions
that anti-unify e-classes and e-nodes.
Note that e-node AU is always invoked in a non-empty context.
The first equation anti-unifies two e-nodes with the same constructor:
in this case, we recursively anti-unify their child e-classes
and return the cross-product of the results.
The second equation applies to e-nodes with different constructors:
as in term AU, this results in a pattern variable.
A nice side-effect of dealing with e-graphs
is that we need not keep track of the anti-substitution
to guarantee that each pair of subterms maps to the same variable:
because any duplicate terms are represented by the same e-class,
we can simply use the e-class ids $a$ and $b$ in the name of the pattern variable $X_{a,b}$.

The second block of equations defines anti-unification of e-classes
(let us first ignore the \textcolor{darkblue}{\sf dominant} which will be explained shortly).
The first equation applies when $a$ and $b$ have already been visited:
in this case, we break the cycle and return the empty set.
Otherwise, the last equation anti-unifies all pairs of e-nodes from the two e-classes
in an updated context and merges the results.
Note that this will add a pattern variable unless all e-nodes in both e-classes
have the same constructor;
we have omitted this detail in \autoref{sec:overview} for simplicity,
but this is implemented in \tname{babble} and often yields more optimal patterns.
For example, consider anti-unifying $c_1$ and $c_2$ from \autoref{fig:egraph} (right):
although they have the constructor \T{scale} in common,
$X$ is actually a better pattern for abstracting these two e-classes than $\T{scale}\  X\ Y$,
because the total size of the actual arguments to pattern $X$ ($\pSide{6}$ and $\scale{\repRot{\pSide{8}}{8}{2\pi/8}}{2}$)
is the same as those to the pattern $\T{scale}\  X\ Y$ ($\pSide{6}$, $1$, $\repRot{\pSide{8}}{8}{2\pi/8}$, and $2$),
and the pattern $X$ itself is smaller.
This happens because the class $c_1$ represents several different terms,
and the term ``compatible with'' $X$ in this case is smaller than the term ``compatible wit'' $\T{scale}\  X\ Y$.

\mypara{Dominant Patterns}
Recall that $\mathsf{AU}(a,b)$ produces patterns
with variable names $X_{a_i,b_i}$, which record the e-class ids they abstract;
let us refer to such a pattern $p$ as \emph{uniquely matched}
and define $\theta_l(p) = \{\many{X_{a_i,b_i} \mapsto a_i}\}$ and $\theta_r(p) = \{\many{X_{a_i,b_i} \mapsto b_i}\}$;
these substitutions are necessarily among $\matches{a}{p}$ and $\matches{b}{p}$, respectively.
Given two uniquely matched patterns $p_1,p_2 \in \mathsf{AU}(a,b)$,
we say that $p_1$ \emph{dominates} $p_2$ in the context of $(a,b)$ if
  (1) $\mathsf{vars}(p_1) \subseteq \mathsf{vars}(p_2)$, and
  (2) $\mathsf{size}(p_1) \leq \mathsf{size}(p_2)$.
We can show that if $p_1$ dominates $p_2$,
then we can safely discard $p_2$ from the set of candidate patterns.
First, since $p_1$ is no larger than $p_2$,
the definition of its $\lambda$-abstraction is also no larger.
Second, given a term $t_a\in \denot{a}$,
compressing this term using $p_1$ vs $p_2$,
requires choosing actual arguments from $\rng{\theta_l(p_1)}$ vs $\rng{\theta_l(p_2)}$;
because the former is a subset of the latter,
the first application can always be made no larger
(symmetric argument applies for $t_b\in \denot{b}$).

Hence it is sufficient that $\mathsf{AU}(a,b)$ only returns the set of \emph{dominant patterns}
(\emph{i.e.}\xspace a pattern dominated by any pattern in the set can be discarded).
This is what the function $\textcolor{darkblue}{\mathsf{dominant}}$ does in the last equation of \autoref{fig:ecau}.
This pruning technique is especially helpful in the presence of equational theories.
Suppose our theory contains the equation $X + Y \equiv Y + X$,
and that the original term $t$ contains subterms $1 + 2$ and $3 + 1$.
After saturation, the e-graph will represent $1 + 2$ and $2 + 1$ in some e-class $a$
and $3 + 1$ and $1 + 3$ in another e-class $b$;
$\mathsf{AU}(a,b)$ will then produce both patterns $X + 1$ and $1 + X$,
but it is clearly redundant to have both,
since they match the same e-classes with the same substitutions $\theta$.
Pruning of dominated patterns will eliminate one of them.

\mypara{Avoiding Cycles}
Interestingly enough,
the same notion of dominant patterns justifies why we need not follow cycles in the e-graph when computing $\mathsf{AU}(a,b)$,
or, alternatively, why a finite set of candidate patterns is sufficient
to compress any term in $\denot{\mathcal{G}}$, even if this set is infinite.
Removing the first equation that short-circuits cycles can only lead to solutions $p'$ of the form
$$
p' = \sapp{[X \mapsto p]}{q}
$$
where $p \in \mathsf{AU}(a,b)$ is another solution, and $q$ is some context, added by the cycle.
It is clear that any such $p'$ is dominated by $p$, and hence can be discarded.

\mypara{E-Graph Anti-Unification}
The algorithm $\mathsf{AU}(a,b)$ computes a set of patterns that can be used to compress terms
represented by the e-classes $a$ and $b$.
Our ultimate goal, however, is to compute candidate patterns for abstracting
\emph{all subterms} of some $t \in \denot{\mathcal{G}}$.
The most straightforward way to achieve this is to apply $\mathsf{AU}(a,b)$ to all pairs of e-classes in $\mathcal{G}$.
We can do better, however:
some pairs of e-classes need not be considered,
because they cannot occur together in a single term $t$.
For an example, consider the following e-graph, with $\mathcal{I} = \mathbb{N}$ and $r = 0$:
$$
0 \mapsto \{f(1),g(2)\}\quad 1\mapsto\{g(3)\} \quad 2\mapsto \{f(3)\} \quad 3\mapsto \{a\}
$$
This e-graph can result, for example, by rewriting a term $f(g(a))$
using an equation $f(g(X)) \equiv g(f(X))$.
In this e-graph, the e-classes $1$ and $2$ (representing $g(a)$ and $f(a)$, respectively)
clearly cannot \emph{co-occur} in the same term:
since the e-graph only represents two terms, $f(g(a))$ and $g(f(a))$.

To formalize this intuition, we define the co-occurrence relation between two e-class ids as follows.
An e-class $a$ is a \emph{sibling} of $b$ if there is an e-node that has both $a$ and $b$ as children.
An e-class $a$ is an \emph{ancestor} of $b$ if $a = b$ or $b$ is a child of some e-node $n \in M(c)$ and $a$ is an ancestor of $c$;
$a$ is a \emph{proper ancestor} of $b$ if $a$ is an ancestor of $b$ and $a \neq b$.
Two e-classes $a$ and $b$ are \emph{co-occurring} if
  (1) one of them is a proper ancestor of another, or
  (2) they have ancestors that are siblings.

Finally, to compute the set $\mathsf{AU}(\mathcal{G})$ of all candidate patterns for an e-graph $\mathcal{G}$,
\T{LLMT} first computes the co-occurrence relation between all e-classes in $\mathcal{G}$,
and then computes $\mathsf{AU}(a,b)$ of all pairs of e-classes that are co-occurring.
As mentioned in \autoref{sec:overview},
we use a dynamic programming algorithm that memoizes the results of $\mathsf{AU}(a,b)$
to avoid recomputation.

\section{Evaluation}
\label{sec:eval}

We evaluated \tname{babble} and the LLMT algorithm behind it
 with two quantitative research questions and a third qualitative one:

\begin{enumerate}[label=RQ \arabic*.,ref=RQ \arabic*]
  \item Can \tname{babble} compress programs better than a state-of-the-art library learning tool?\label{rq:1}
  \item Are the main techniques in LLMT (anti-unification and equational theories)
        important to the algorithm's performance?\label{rq:2}
  \item Do the functions \tname{babble} learns make intuitive sense?\label{rq:3}
\end{enumerate}

\mypara{Benchmark Selection}
We use two suites of benchmarks to evaluate \tname{babble}, both shown in \autoref{tab:benchmarks}.
The first suite originates from the \tname{DreamCoder} work~\cite{dreamcoder},
 and is available as a public repository~\cite{compression-bench}.
\tname{DreamCoder} is the current state-of-the-art library learning tool,
and using these benchmarks allows us to perform a head-to-head comparison.
The \tname{DreamCoder} benchmarks are split into five domains (each with a different DSL);
we selected two of the domains---List and Physics---which we understood best,
to add an equational theory.

The second benchmark suite, called \tname{2d cad}, comes from \citet{cogsci-dataset}.
This work collects a large suite of programs in a graphics DSL for the
 purpose of studying connections between the generated objects and their
 natural language descriptions.
There are 1,000 programs in the ``Drawings'' portion of this dataset,
 divided into four subdomains (listed in \autoref{tab:benchmarks})
 of 250 programs each.\looseness=-1

We ran \tname{babble} on all benchmarks on an AMD EPYC 7702P processor at 2.0 GHz.
Each benchmark was run on a single core.
The \tname{DreamCoder} results were taken from
 the benchmark repository~\cite{compression-bench};
 \tname{DreamCoder} was run on 8 cores of an AMD EPYC 7302 processor at 3.0 GHz.

\begin{table}
  \centering
  \begin{tabular}{lrr}
    \multicolumn{3}{c}{\tname{DreamCoder}~\cite{dreamcoder}} \\
    \bf Domain & \bf \# Benchmarks & \bf \# Eqs \\
    List    & 59 & 14 \\
    Physics & 18 & 8 \\
    Text    & 66 & - \\
    Logo    & 12 & - \\
    Towers  & 18 & -
  \end{tabular}\qquad%
  \begin{tabular}{lrr}
    \multicolumn{3}{c}{\tname{2d cad}~\cite{cogsci-dataset}} \\
    \bf Domain & \bf \# Benchmarks & \bf \# Eqs \\
    Nuts \& Bolts & 1 & 7 \\
    Vehicles      & 1 & 9 \\
    Gadgets       & 1 & 17 \\
    Furniture     & 1 & 9 \\
    &
  \end{tabular}
  \caption{
    We selected our benchmark domains from two previous works:
     \tname{DreamCoder}~\cite{dreamcoder}
     and \tname{2d cad}~\cite{cogsci-dataset}.
    Each domain from \tname{DreamCoder} has multiple benchmarks;
     \tname{2d cad} has one large benchmark per domain.
    For some domains,
     we additionally supplied \tname{babble} with an equational theory;
     we report the number of equations in the final column.
  }\label{tab:benchmarks}
\end{table}

\subsection{Comparison with \tname{DreamCoder}}

\begin{figure}
  \centering
  \begin{subfigure}{0.48\linewidth}
    \includegraphics[width=\linewidth]{figures/scatter/list.pdf}
    \caption{List domain}
    \label{fig:dc-domains-list}
  \end{subfigure}
  \hfill
  \begin{subfigure}{0.48\linewidth}
    \includegraphics[width=\linewidth]{figures/scatter/physics.pdf}
    \caption{Physics domain}
    \label{fig:dc-domains-physics}
  \end{subfigure}

  \begin{subfigure}{0.48\linewidth}
    \includegraphics[width=\linewidth]{figures/scatter/text.pdf}
    \caption{Text domain}
  \end{subfigure}
  \hfill
  \begin{subfigure}{0.48\linewidth}
    \includegraphics[width=\linewidth]{figures/scatter/logo.pdf}
    \caption{Logo domain}
  \end{subfigure}

  \begin{subfigure}{0.48\linewidth}
    \includegraphics[width=\linewidth]{figures/scatter/towers.pdf}
    \caption{Towers domain}
  \end{subfigure}
  \caption{
    \tname{babble} consistently achieves better compression ratios than
     \tname{DreamCoder} on benchmarks from the \tname{DreamCoder} domains,
     and it does so 1--2 orders of magnitude faster.
    Each marker shows the compression ratio (x-axis)
     and run time (y-axis) of a benchmark.
    Each benchmark is one \tname{DreamCoder} input,
     \emph{i.e.}\xspace, a set of groups of programs as described above.
    Lower and to the right is better.
    In the domains where we supplied
     \tname{babble} with an equational theory
     (List and Physics),
     additional markers show the performance
     of \tname{babble} using purely syntactic learning (without equations, ``BabbleSyn'')
     or only equality saturation without library learning (``EqSat'').
  }\label{fig:dc-domains}
\end{figure}

To answer \ref{rq:1},
 we compare to the state-of-the-art \tname{DreamCoder} tool~\cite{dreamcoder} on its own benchmarks.
The \tname{DreamCoder} benchmarks are suited to its workflow;
while the input to a library learning task
 is conceptually a set of programs (or just one big program),
 each input to \tname{DreamCoder} is a set of \emph{groups} of programs.
Each group is a set of programs
 that are all output from the same program synthesis task
 (from an earlier part of the \tname{DreamCoder} pipeline).
When compressing a program via library learning,
 \tname{DreamCoder} is minimizing the cost of the program
 made by concatenating the most compressed
 program from each group together, in other words:
\[
  \sum_{\textsf{group $g$}}
  \min_{\textsf{program $p \in g$}}
  \textsf{cost}(p)
\]
\tname{DreamCoder} takes this approach to give its library learning
 component many variants of the same program,
 in order to introduce more shared structure
 between solution programs across different synthesis problems.

To implement \tname{DreamCoder}'s benchmarks in \tname{babble},
 we use the e-graph\xspace to capture the notion of program variants in a group.
Since every program in a group is the output of the same synthesis task,
 \tname{babble} considers them equivalent and places them in the same e-class\xspace.

\mypara{Results}
We ran \tname{babble} on five domains from the \tname{DreamCoder} benchmark suite.
The results are shown in \autoref{fig:dc-domains}.
In summary,
 \tname{babble} consistently achieves better compression ratios than
 \tname{DreamCoder} on benchmarks from the \tname{DreamCoder} domains,
 and it does so 1--2 orders of magnitude faster.

\mypara{The Role of Equational Theory}
\label{sec:ablation}
To answer \ref{rq:2},
 we again turn to the \tname{DreamCoder} benchmarks,
 focusing on the domains where we supplied \tname{babble}
 with an equational theory: List and Physics.
In these domains, we ran \tname{babble} in two additional configurations:
\begin{itemize}
  \item ``BabbleSyn'' ignores the equational theory, just doing syntactic library learning.
  \item ``EqSat'' just optimizes the program using Equality Saturation with the
        rewrites from the equational theory.
        This configuration \emph{does not} do any library learning.
\end{itemize}
\autoref{fig:dc-domains-list} and \ref{fig:dc-domains-physics}
 show the results for these additional configurations,
 as well as \tname{DreamCoder} and the normal \tname{babble} configuration.
All \tname{babble} configurations rely on
  targeted subexpression elimination
  to select the final learned library.
The ``EqSat'' configuration is unsurprisingly very fast
 but performs relatively little compression,
 as it does not learn any library abstractions.
The ``BabbleSyn'' configuration
 does indeed compress the inputs,
 in fact it is still better than \tname{DreamCoder} in both domains.
However,
 the addition of the equational theory (the ``\tname{babble}'' markers in the plots)
 significantly improves compression
 and adds relatively little run time,
 well within an order of magnitude.

\subsection{Large-Scale \tname{2d cad} Benchmarks}

\begin{table}
  \begin{tabular}{lr|rrr|rrr}
    \multicolumn{2}{c}{} & \multicolumn{3}{c}{Without Eqs} & \multicolumn{3}{c}{With Eqs} \\
    \bf Benchmark & \bf Input Size & \bf Out Size & \bf CR & \bf Time (s) & \bf Out Size & \bf CR & \bf Time (s) \\
    Nuts \& Bolts   &      19009 &       2059 &       9.23 &      18.74 &       1744 &      10.90 &      40.75 \\
    Vehicles        &      35427 &       6477 &       5.47 &      79.50 &       5505 &       6.44 &      78.03 \\
    Gadgets         &      35713 &       6798 &       5.25 &      75.07 &       5037 &       7.09 &      82.29 \\
    Furniture       &      42936 &      10539 &       4.07 &     133.25 &       9417 &       4.56 &     110.00 \\
    \hline
    Nuts \& Bolts (clean) & 18259 &       2215 &       8.24 &      18.12 &       1744 &      10.47 &      40.91
  \end{tabular}
  \vspace{1em}
  \captionof{table}{
    Results for running \tname{babble} on the four large examples
     from the \tname{2d cad} dataset~\cite{cogsci-dataset},
     both without and with an equational theory.
    Each row includes the input and output AST sizes,
     the compression ratio (CR),
     and \tname{babble}'s run time in seconds.
    \autoref{fig:cogsci-scatter} plots this data.
    The final row shows performance on a modified dataset.
  }\label{tab:cogsci}
\end{table}
\begin{figure}
  \includegraphics[height=5cm]{figures/scatter/cogsci-scatter.pdf}
  \caption{
    Data from \autoref{tab:cogsci} in scatter plot.
    Each line segment shows compression ratio and run time for a domain
     without (hollow marker) and with (solid marker) an equational theory.
    Using an equational theory improves compression in all cases,
     and even improves run time in two cases.
  }\label{fig:cogsci-scatter}
\end{figure}

The previous section demonstrated that \tname{babble}'s performance
 far surpasses the state of the art.
In this section,
 we present and discuss the results
 of running \tname{babble} on benchmarks from the \tname{2d cad} domain.
These benchmarks,
 taken from \citet{cogsci-dataset},
 are significantly larger (roughly 10x - 100x)
 than those from the \tname{DreamCoder} dataset
 and out of reach for \tname{DreamCoder}.

\mypara{Quantitative Results}
\autoref{tab:cogsci} and \autoref{fig:cogsci-scatter} show the results
 of running \tname{babble} on the benchmarks from the \tname{2d cad} domain.
The plot in \autoref{fig:cogsci-scatter} makes two observations clear.
First and relevant for \ref{rq:2},
  the addition of an equational theory improved all four benchmarks
  (the solid marker is always to the right of the hollow marker).
Second, and perhaps surprisingly,
 equational theories can sometimes make \tname{babble} \emph{faster}!
This is consistent with previous observations about equality saturation~\cite{willsey2021egg}:
 while equality saturation typically makes an e-graph larger,
 it can sometimes combine two relevant e-classes into one and
 reduce the amount of work that some operation over an e-graph must do.

We also observed that the Nuts-and-bolts\xspace dataset contains
  several redundant transformations,
  like the ``scale by 1'' featured in the running example of \autoref{sec:overview}.
These redundancies can be useful for finding
  abstractions in the absence of an equational theory.
However, they should not be required in
  \tname{babble} since LLMT can
  introduce the redundancies wherever
  required.
We therefore removed all
  existing redundant transformations
  from Nuts-and-bolts\xspace and ran \tname{babble}
  on the transformed dataset.
The results are in the final row of \autoref{tab:cogsci}.
On the modified dataset,
 \tname{babble} achieves identical compression when using the equational theory,
 but without the equations it performs worse than on the unmodified dataset.

\mypara{Qualitative Evaluation}
\label{sec:quali}
\begin{figure}
  \centering
  \includegraphics[width=\textwidth]{figures/qual_cogsci.pdf}
  \caption{A selection of evaluated programs from each of the domains
           in the \tname{2d cad} dataset, along with a selection of functions
           that \tname{babble} learns within the first ten rounds for each domain.
           Bolded functions represent learned abstractions.
           Note this figure uses the concrete syntax from the \tname{2d cad}
           dataset; it is similar to the simplified form shown in the overview.
           We named the learned functions and their parameters for clarity.
   }\label{fig:qual_cogsci}
\end{figure}
\autoref{fig:qual_cogsci}
 highlights a sample of abstractions that
 \tname{babble} discovered from the \tname{2d cad} benchmarks.
We ran \tname{babble} on each of the benchmarks
  and applied the learned abstractions on a few input programs to
  visualize their usage.
Questions about usability and readability of
  learned libraries are difficult to answer
  without rigorous user studies which we leave for
  future work.
Nevertheless, \autoref{fig:qual_cogsci} shows that
  \tname{babble} identifies common structures that
  are similar across
  different benchmarks, which
  makes its output easier to reuse and interpret.

First, we revisit the Nuts-and-bolts\xspace example
  from \autoref{sec:intro}:
  \autoref{fig:qual_cogsci} shows that \tname{babble}
  learns the scaled polygon (\texttt{ngon}) abstraction which
  is applicable to several programs in the dataset.
We also see that \tname{babble} consistently
  finds a similar abstraction
  representing a ``ring of shapes'' for both
  Nuts-and-bolts\xspace and Vehicles.
Finally, as the Gadgets example shows,
  \tname{babble} finds abstractions for both
  the entire model as well as its components.
In this case,
  it learned the function \texttt{gadget\_body}
  that abstracts the entire outer shape,
  and it also learned \texttt{dial} that abstracts
  the handles of the outer shape.\looseness=-1

\section{Introduction}
\label{sec:intro}

Abstraction is the key to managing software complexity.
Experienced programmers routinely extract common functionality
into libraries of reusable abstractions
to express their intent more clearly and concisely.
What if this process of extracting useful abstractions from code could be automated?
\emph{Library learning} seeks to answer this question
  with techniques to compress
  a given corpus of programs
  by extracting common structure into reusable library functions.
Library learning
  has many potential applications
  from refactoring and decompilation~\cite{szalinski,shapemod},
  to modeling human cognition~\cite{cogsci-smiley,cogsci-dataset},
  and speeding up program synthesis by
  specializing the target language to
  a chosen problem domain~\cite{dreamcoder}.

Consider the simple library learning task in \autoref{fig:polygons}.
On the left,
 \autoref{fig:polygons}a shows a corpus of three programs in a
 \tname{2d cad} DSL from \citet{cogsci-dataset}.
Each program corresponds to a picture
 composed of regular polygons,
 each of which is made of multiple rotated line segments.
On the right,
 \autoref{fig:polygons}b shows a learned library
 with a single function (named \T{f0})
 that abstracts away the construction of scaled regular polygons.
The three input programs can then be refactored
 into a more concise form using the learned \T{f0}.
Whether \T{f0} is the ``best'' abstraction for this corpus is generally hard to quantify.
For this paper,
 we follow \tname{DreamCoder}~\cite{dreamcoder} and use \emph{compression} as a
 metric for library learning, \emph{i.e.}\xspace,
 the goal is to reduce the total size of the corpus in AST nodes
 (from 208 to 72 \autoref{fig:polygons}).
Importantly, the total size of the corpus includes the size of the library:
this prevents library learning from generating too many overly-specific functions,
and instead biases it towards more general and reusable abstractions.

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/polygons.pdf}
    \caption{Example of library learning.
             Initial corpus A contains three graphical programs from the ``nuts \& bolts'' dataset of \citet{cogsci-dataset}.
             Corpus B is the output of library learning with a single learned function for a scaled polygon,
             and the original programs refactored using this function.}\label{fig:polygons}
\end{figure}

Library learning can be phrased as a program synthesis problem
 structured in two phases:
\emph{generating} candidate abstractions, and then
\emph{selecting} those abstractions that produce the best (smallest)
     refactored corpus.
The state-of-the-art technique, implemented in \tname{DreamCoder}~\cite{dreamcoder},
  suffers from two primary limitations
  that hinder scaling library learning to larger and more realistic inputs.
\begin{itemize}
    \item
    Candidate generation is not \emphbf{precise}:
    \tname{DreamCoder} generates many candidate abstractions that cannot be useful,
    slowing down the selection phase and the algorithm as a whole.

    \item

    The technique is purely syntactic and hence not \emphbf{robust} to superficial variation.

    For example,
     a human programmer would immediately know that the terms $2 + 1$ and $1 + 3$
     can be refactored using the abstraction $\lambda x \ \mathbf{\shortrightarrow}\  x + 1$,
     because addition commutes;
     a purely syntactic library learning approach cannot generate this abstraction.

\end{itemize}
In this paper we propose \emph{library learning modulo (equational) theories} (LLMT)---%
a new library learning algorithm
that addresses both of the above limitations.

\mypara{Precise Candidate Generation via Anti-Unification}
To make candidate generation more precise,
 LLMT leverages two key observations:
\begin{itemize}
    \item
    Useful abstractions must be used in the corpus at least twice.
    %
    For example, in a corpus of two programs $2 + 1$ and $3 + 1$,
    there is no need to consider $\lambda x \ \mathbf{\shortrightarrow}\  3 + x$
    because it can only be used in one place,
    and hence would only increase the size of the corpus.

    \item
    Abstractions should be ``as concrete as possible'' for a given corpus.
    For example,
     in the same corpus with $2 + 1$ and $3 + 1$,
     the abstraction $\lambda x \ \mathbf{\shortrightarrow}\  x + 1$
     is superior to the more general
     $\lambda x\ y \ \mathbf{\shortrightarrow}\  x + y$,
     since both apply to the same two terms,
     but applying the latter is more expensive (it requires two arguments).

\end{itemize}
In other words, a useful abstraction corresponds to the least general \emph{pattern}
that matches some pair of subterms from the original corpus;
such a pattern can be computed via \emph{anti-unification} (AU)~\cite{plotkin1970lattice}.
For example, anti-unifying $2 + 1$ and $3 + 1$
  yields the pattern $X + 1$,
  and the desired candidate library function $\lambda x \ \mathbf{\shortrightarrow}\  x + 1$
  can be derived by abstracting over the pattern variable $X$.

Similarly, in \autoref{fig:polygons}, the abstraction \T{f0} can be derived
 by anti-unifying, for example, the blue and the brown subterms of the corpus.

\mypara{Robustness via E-Graphs}
To make candidate generation more robust,
LLMT additionally takes as input a \emph{domain-specific equational theory}
and uses it to find programs that are semantically equivalent to the original corpus,
but share the most syntactic structure.
For example, in the domain of arithmetic, we expect the theory to contain the equation
$X + Y  \equiv  Y + X$,
which states that addition is commutative.
Given the corpus with terms $2 + 1$ and $1 + 3$,
this rule can be used to rewrite the second term in to $3 + 1$,
enabling the discovery of the common abstraction $\lambda x \ \mathbf{\shortrightarrow}\  x + 1$.

The main challenge with this approach is to search over the large space of programs
that are semantically equivalent to the original corpus.
To this end, LLMT relies on the \emph{e-graph} data structure
and the \emph{equality saturation} technique~\cite{tate2009equality,willsey2021egg}
to compute and represent the space of semantically equivalent programs.
To enable efficient library learning over this space,
we propose a new candidate generation algorithm
that efficiently computes the set of all anti-unifiers
between pairs of sub-terms represented by an e-graph,
using dynamic programming.

Finally, \emph{selecting} the optimal library
  in this setting reduces to the problem of
  \emph{extracting} the smallest term out of an e-graph
  in the presence of common sub-expressions.
This problem is extremely computationally intensive in its general form,
and existing approaches have limited scalability~\cite{tensat}.
Instead we propose \emph{targeted subexpression elimination}:
a novel e-graph extraction algorithm
that uses domain-specific knowledge to reduce the search space
and readily supports approximation via beam search to trade off accuracy and efficiency.

\begin{figure}
    \centering
    \includegraphics[width=.8\textwidth]{figures/arch.pdf}
    \caption{\tname{babble} architecture overview}\label{fig:arch}
  \end{figure}

\mypara{\tname{babble}}
We have implemented LLMT in \tname{babble},
a tool built on top of the \tname{egg} e-graph library~\cite{willsey2021egg}.
The overview of the \tname{babble}'s workflow is shown in \autoref{fig:arch},
with gray boxes representing existing techniques
and black boxes representing our contributions.

We evaluated \tname{babble} on benchmarks from two sources:
compression tasks extracted from \tname{DreamCoder} runs~\cite{compression-bench}
and \tname{2d cad} programs used to evaluate concept learning in humans~\cite{cogsci-dataset}.
Our evaluation shows that \tname{babble} outperforms \tname{DreamCoder} on its own benchmarks,
 achieving better compression in orders of magnitude less time.
Adding domain-specific rewrites improves compression even further.
We also show that \tname{babble} scales to the larger \tname{2d cad} corpora,
which is beyond reach of \tname{DreamCoder}.
We also present and discuss a selection of abstractions learned by \tname{babble},
 demonstrating that the LLMT approach can learn functions that match human intuition.

\mypara{Contributions}
In summary, this paper makes the following contributions:
\begin{itemize}
    \item \emph{library learning modulo equational theory}:
        a library learning algorithm that can incorporate an arbitrary domain-specific equational theory
        to make learning robust to syntactic variations;
    \item \emph{e-graph anti-unification}:
        an algorithm that efficiently generates the set of candidate abstractions
        using the mechanism of anti-unification extended from terms to e-graphs;
  \item \emph{targeted common subexpression elimination}:
        a new approximate algorithm for extracting the best term from an e-graph
        in the presence of common sub-expressions.
\end{itemize}

\begin{comment}

In this work, we consider the problem of library learning
as formulated in \tname{DreamCoder}~\cite{dreamcoder}.
At a high level, the goal of library learning is to \emph{compress} a given corpus of programs
by inventing a library of abstractions and refactoring the corpus using this library.
\autoref{fig:polygons} shows a simple instance of library learning
over three graphical programs sampled from the dataset of \citet{cogsci-dataset}.
The original programs are expressed in a vector graphics DSL
with constructs for primitive shapes (\emph{e.g.}\xspace \T{line}),
applying a transformation matrix to a shape (\T{trans}),
and repeating a shape $n$ times, with a transformation applied between iterations (\T{repeat}).
Although this is not immediately obvious from the programs themselves,
looking at the corresponding images makes it clear that
each program describes several nested polygons (hexagons and octagons) at different scales.
If we abstract the concept of a scaled polygon into a two-argument function $f_0$,
we can refactor the original programs to be much more concise, as shown on the right-hand side of \autoref{fig:polygons},
shrinking the total size of the corpus (including the library) from 208 to 72 AST nodes,
for a compression rate of $\sim 2.9$.
This is the best compression we can achieve with a single-function library,%
\footnote{We can achieve better compression if we allow a second function $f_1$,
which specializes $f_0$ to the case of a hexagon;
we focus on the one-function library here for simplicity of presentation.}
and hence this is the desirable outcome of library learning.
\TODO{I'm torn here:
on the one hand, it would be nice to show a multi-function library,
but on the other hand, with 2 functions, the difference between w and w/o DSRs becomes super small and not impressive
(the corpus is so small, it's very easy to overfit).}

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/polygons.pdf}
    \caption{Example of library learning.
             Initial corpus A contains three graphical programs from the ``nuts \& bolts'' dataset of \citet{cogsci-dataset}.
             Corpus B is the output of library learning with a single learned function for a scaled polygon,
             and the original programs refactored using this function.}\label{fig:polygons}
\end{figure}

\mypara{Limitations of Prior Work}
\tname{DreamCoder}'s approach to library learning, called \emph{beta-inversion},
discovers library functions one at a time,
alternating between two stages:
\begin{enumerate*}[(1)]
    \item proposing candidate abstractions and
    \item selecting the optimal abstraction (the one that results in the best compression).
\end{enumerate*}
The first limitation of this approach is that the proposal stage is completely \emph{unguided}:
beta-inversion simply picks a random program fragment from the initial corpus
and replaces some of its arbitrarily chosen subterms with abstraction variables.
For example, to learn the function $f_0$ in \autoref{fig:polygons},
\tname{DreamCoder} would have to be lucky enough to pick, say, the orange subterm as a fragment,
\emph{and then also} decide to abstract over its subterms $6$ and $2.25$.
As a result, this approach does not scale beyond small corpora of short programs.

The second limitation beta-inversion
is that is can only extract purely \emph{syntactic} similarities between program fragments,
and hence it is \emph{not robust} to semantics-preserving variations in the input corpus.
For an example, let us take a closer look at \autoref{fig:polygons}.
Note that the terms for the two blue hexagons are of the form $\T{trans}\ (\T{repeat}\ \ldots)\ [1\ 0\ 0\ 0]$,
which amounts to scaling the underlying unit hexagon by 1!
This redundant transformation is crucial
for a syntactic library learning approach to generate the optimal abstraction $f_0$ from \autoref{fig:polygons}.
As shown in \autoref{fig:llmt},
eliminating this transformation
causes beta-inversion to learn a different function, $f_1$,
which abstracts over a unit (unscaled) polygon.
Although the simplified input corpus C from \autoref{fig:llmt}
is smaller than the original corpus A (196 AST nodes instead of 208),
its compressed version D is larger  (81 AST nodes instead of 72),
because the scaling transformation is now outside the abstraction and has to be repeated five times.

In this paper we propose \emph{library learning modulo (equational) theories} (LLMT)---%
a new library learning algorithm,
which addresses both of the above limitations.

\mypara{Scalability with Anti-Unification}
To make the proposal stage of library learning more efficient,
LLMT leverages two key observations:
\begin{enumerate*}[(1)]
    \item In order to be compressive, an abstraction must appear in the corpus at least twice;
    for example, there is no need to consider $\lambda x \ \mathbf{\shortrightarrow}\  \T{trans}\ x\ [0.5\ 0\ 0\ 0]$ in our running example,
    because it can only be used in one place, thereby only increasing the size of the corpus.
    \item An abstraction cannot be optimal unless it captures all the common structure between the terms it abstracts;
    for example, there is no need to consider $\lambda x\ y to \T{repeat}\ (\T{trans line}\ \ldots)\ x\ y$ either,
    because everywhere this pattern occurs, $y$ actually has the shape $[1\ (2\pi / x)\ 0\ 0]$.
\end{enumerate*}
In other words, any viable candidate abstraction is a \emph{least general generalization} of some pair of subterms from the original corpus,
and can be computed via \emph{anti-unification} (AU)~\cite{plotkin1970lattice}.
For example, the abstraction $f_0$ can be computed by anti-unifying the green and the brown subterms of the corpus in \autoref{fig:polygons}.
The remaining challenge is to efficiently compute the set of anti-unifiers
between \emph{all} pairs of subterms in the corpus;
LLMT achieves this via a new bottom-up dynamic programming algorithm.

\begin{figure}
    \centering
    \includegraphics[width=\textwidth]{figures/polygons-simplified.pdf}
    \caption{Difference between syntactic library learning and LLMT.
             Here the initial corpus C is the simplified version of corpus A from \autoref{fig:polygons},
             with the redundant transformations on the blue subterms removed.
             With this modification, syntactic techniques would learn a suboptimal abstraction $f_1$,
             leading to corpus D,
             while LLMT still learns the optimal abstraction $f_0$.
     }\label{fig:llmt}
\end{figure}

\mypara{Robustness with E-Graphs}
To address the robustness issue,
LLMT additionally takes as input a domain-specific equational theory (expressed as a set of rewrite rules)
and uses it to find programs that are semantically equivalent to the original corpus,
but share the most syntactic structure.
For example, for the graphical DSL from our running example,
the equational theory might contain a rule that adds a unit transformation to any shape $s$:
$$
s  \Rightarrow  \T{trans}\ s\ [1\ 0\ 0\ 0]
$$
This rule can then be used to transform the two unit hexagons in corpus C from \autoref{fig:llmt}
into the form they have in the original corpus in \autoref{fig:polygons},
which enabled the discovery of the scaled polygon abstraction $f_0$
and using it to compress the corpus to size 72.

The main challenge with this approach is how do we represent and learn over the large space of programs
that are semantically equivalent to the original corpus.
To this end, LLMT relies on the \emph{e-graph} data structure
and the \emph{equality saturation} technique~\cite{tate2009equality,willsey2021egg}
to compute and represent the space of semantically equivalent programs.
To enable efficient library learning over this space,
we extend our AU-based algorithm for suggesting candidate abstractions to work over e-graphs.
We also propose a new algorithm for
  selecting the best subset of candidate abstractions,
  based on \emph{targeted common subexpression elimination}
  extraction for e-graphs.
\TODO{Do we need to talk about this? If we end up learning one at a time,
our selection algorithm is not that different from DC.}.
\TODO{ztatlock: I think it is ok to keep this?}

\end{comment}

\section{Old Intro}

When creating software, programmers normally choose not to place their
entire program in a single function; instead, programmers will often abstract
out common functionality into helper functions, which can be reused and composed
to make more complex functionality.
\emph{Library learning} is the problem of how to automate this process: creating
an algorithmic approach which, when given an input set of programs, can
identify common structure and learn new functions/abstractions.
The most obvious use-case for a tool which could tackle this problem is
for automated code refactoring: automatically learning a set of functions
in order to reduce code size as much as possible.
One prior example of this is Szalinski, which rewrites programs representing
CAD models with the goal of compressing them and making them easily editable
later on \cite{szalinski}.

However, library learning also has many other use cases,
several of which have been documented in prior literature.
One possible use case for library learning is as a component in a larger
program synthesis loop.
When faced with the task of synthesizing a program which can correctly solve
a task, many program synthesis approaches rely on having a set of known
library functions which can be composed to create solution programs.
This set of library functions is normally hard-coded by the programmer.
However, with the addition of a library learning algorithm, a program synthesis
procedure would be able to first synthesize solutions to a set of tasks,
then learn common functions which are used across those solutions,
and finally use those new common functions in synthesizing solutions for other
tasks.
By introducing an intermediate library learning step in between synthesis
stages, the system can learn fundamental operations like map and fold and
synthesize progressively more complex programs, without user intervention.
This approach to program synthesis was pioneered by DreamCoder,
which demonstrated that this iterative method was able to learn these
fundamental abstractions and use them to synthesize complex programs
which could solve a variety of tasks across several different domains
\cite{dreamcoder}.

Another use case for algorithmic library learning is to model how humans can
intuitively recognize abstract visual common structure.
For instance, humans are naturally able to recognize smiley faces, regardless
of the actual level of detail of the smiley faces themselves.
Figuring out how to get computers to solve the same task and recognize common
visual structure can serve as a useful proxy to model how humans can
instinctively accomplish the same task.
This cognitive science-focused use case for library learning has been explored
in \cite{cogsci-smiley}.

\mypara{Shortcomings of the existing state-of-the-art}
As we've shown, finding a general, efficient, and effective approach to
library learning can have ramifications across a wide variety of different
domains.
However, \TODO{existing approaches can't do this}.
Of the prior work discussed, DreamCoder's library learning approach is the
most general and can apply to a wide variety of domains.
This approach relies on \TODO{``inverse-beta reduction'':
essentially poking random holes in programs to guess-and-check where
function abstractions should go, then picking the best learned function
one-at-a-time}.
\TODO{This approach, while general, doesn't satisfy the efficient or effective
conditions.}
One shortcoming of this approach is its efficiency. One part of this is how
candidate abstractions are \emph{proposed}. Since DreamCoder's library learning
algorithm works by randomly choosing function bodies and randomly choosing
which parts of those expressions should become function arguments, the choice
of what expressions should be abstracted over and where function arguments
should be placed are decided in a random, undirected way, and enumerating over
all possible function bodies and argument places across all input programs would
be prohibitively costly.
\TODO{%
    idk how DreamCoder actually deals with this issue in practice but
    i wanna say somethin like:
    In practice, for efficiency's sake, DreamCoder limits its proposed function
    candidates to some expression size or some number of arguments, which can
    cause it to potentially miss out on useful abstractions.
}
The other part of DreamCoder's efficiency problem is in how it selects the
optimal set of library functions to use to refactor a set of input programs.
In particular, DreamCoder's library learning approach simply iterates over
all abstraction candidates and chooses the best candidate one-at-a-time.
This approach is inefficient but also doesn't take into account how
combining multiple libraries can potentially reduce cost more than
just a set of individually picked library functions.
\TODO{%
    basically i'm tryna say the one-at-a-time approach doesn't take
    into account if you have two libs which, when picked together,
    compress more than if u were to just pick the top two libs which,
    on their own, compress the best%
}

Another shortcoming of this approach is that, while it's a general way of
recognizing common syntactic structure (where different programs look the same
syntactically), it fails at recognizing common semantic structure (where
different programs have similar output behavior).
For instance, in the domain of smiley faces described above, if two programs
produced identical smiley faces with circular eyes, but one program had rotated
the eyes by 90 degrees, DreamCoder's approach would be unable to find the common
smiley face structure between the two programs, since it wouldn't be able to
\TODO{tell} that circles are rotationally symmetric.
\TODO{this para is very handwavy}

\TODO{%
    Discuss the other library learning approaches?
    Are we tryna frame this just as ``improved DreamCoder'', or as
    ``DreamCoder's generality combined with
    others' efficiency and domain-specificity''?%
}
\section{Overview}
\label{sec:overview}

\newcommand{\xargs}[4]{\ensuremath{[#1, #2, #3, #4]}}
\newcommand{\xform}[2]{\ensuremath{\T{xform}(#1,\ #2)}}
\newcommand{\rep}[3]{\ensuremath{\T{repeat}(#1,\ #2,\ #3)}}
\newcommand{\comb}[2]{\ensuremath{\T{combine}(#1,\ #2)}}

\newcommand{\scale}[2]{\ensuremath{\T{scale}(#1,\ #2)}}
\newcommand{\repRot}[3]{\ensuremath{\T{repRot}(#1,\ #2,\ #3)}}
\newcommand{\rotate}[2]{\ensuremath{\T{rotate}(#1,\ #2)}}
\newcommand{\translate}[3]{\ensuremath{\T{translate}(#1,\ #2,\ #3)}}

\newcommand{\pSide}[1]{\ensuremath{\T{side}(#1)}}
\newcommand{\ngon}[1]{\ensuremath{\T{ngon}(#1)}}
\newcommand{\scaledHex}[1]{\ensuremath{\T{scaledHexagon}(#1)}}

We illustrate LLMT via a running example building on \autoref{fig:polygons}.
The input corpus in \autoref{fig:polygons} is written in a \tname{2d cad} DSL by \citet{cogsci-dataset},
which features the following primitives:

\vspace{0.25em}
\begin{tabular}{rl}
  \T{line} &
    a line segment from
    the origin (0, 0) to the point (1, 0)
    \\
  \comb{F_1}{F_2} &
    the union of figures $F_1$ and $F_2$
    \\
  \xform{F}{\tau} &
    applying transformation $\tau$ to figure $F$
    \\
  \rep{F}{0}{\tau} &
    the empty figure
    \\
  \rep{F}{n+1}{\tau} &
    \comb{F}{\xform{\rep{F}{n}{\tau}}{\tau}}, similar to a ``fold''
\end{tabular}
\vspace{0.25em}

A transformation $\tau$ is a 4-tuple
  $\xargs{s}{\theta}{t_x}{t_y}$ denoting
  uniformly scaling by a factor of $s$,
  rotating by $\theta$ radians, and
  translating by $(t_x,\ t_y)$ in the
  $x$ and $y$ directions respectively.
For example, the green hexagon in \autoref{fig:polygons} is implemented as:
\[
  \xform{
    \rep{
      \xform{
        \T{line}
      }{
        \xargs{1}{0}{-0.5}{0.5 / \tan(\pi / 6)}
      }
    }{6}{
      \xargs{1}{2\pi/6}{0}{0}
    }
  }{
    \xargs{2}{0}{0}{0}
  }
\]
That is, a hexagon side $\xform{\T{line}}{\xargs{1}{0}{-0.5}{0.5 / \tan(\pi / 6)}}$
is repeated six times, each time rotated by another $2\pi/6$ radians,
and the resulting unit hexagon is scaled by $2$.

Taking a closer look at
  the two blue hexagons in \autoref{fig:polygons},
  however, we notice something peculiar.
The two occurrences of
\xform{\T{repeat} \ldots}{\xargs{1}{0}{0}{0}}

  in \autoref{fig:polygons} are no-ops:

  they merely scale a figure $F$ by a factor of 1
  and neither rotate nor translate it.
These redundant transformations would likely not be there
had the code been written by hand
or decompiled from a lower-level representation (by a tool like \tname{Szalinski}~\cite{szalinski});%
  \footnote{
    We suspect that these transformations
    ended up in the dataset of~\citet{cogsci-dataset}
    because it was generated programmatically from
    human-designed abstractions, such as ``scaled polygon''.}
and yet, they are crucial
  if we hope to learn the optimal abstraction $f_0$
  with a purely syntactic technique.

\autoref{fig:llmt} shows a simplified and ``more natural''
  version of the corpus from \autoref{fig:polygons},
  which eliminates these no-op transformations.
As illustrated in the figure,
  this causes a purely syntactic technique
  to learn a different function, $f_1$,
  which abstracts over an \emph{unscaled} unit polygon.
Because
  the scaling transformation
  is now outside the abstraction,
  it must be repeated five times.
As a result,
  although the simplified input corpus C from \autoref{fig:llmt}
  is \emph{smaller} than the original corpus A (196 AST nodes instead of 208),
  its compressed version D is \emph{larger} (81 AST nodes instead of 72)!
In other words,
  syntactic library learning is not robust
  with respect to semantics-preserving code variations.

\begin{figure}
  \centering
  \includegraphics[width=\textwidth]{figures/polygons-simplified.pdf}
  \caption{Difference between syntactic library learning and LLMT.
           Here the initial corpus C is the simplified version of corpus A from \autoref{fig:polygons},
           with the redundant transformations on the blue subterms removed
           (unchanged terms are shown as ellipses).
           With this modification, syntactic techniques would learn an inferior abstraction $f_1$,
           leading to corpus D,
           while LLMT still learns the better abstraction $f_0$.
   }\label{fig:llmt}
\end{figure}

In contrast, our tool \tname{babble}
  can take the simplified corpus C as input
  and still discover, in a fraction of a second,
  the scaled polygon abstraction $f_0$,
  yielding the compressed corpus B of size 72.
In the rest of this section,
  we illustrate how \tname{babble} achieves this using our new algorithm,
  \emph{library learning modulo equational theory} (LLMT).

\mypara{Simplified DSL}
In the rest of this section
  we use a tailored version of the \tname{2d cad} DSL
  with the following additional constructs:

  \vspace{0.25em}
  \begin{tabular}{rl}
    \scale{F}{s} &
      scale $F$ by a factor of $s$
      \\
    \repRot{F}{n}{\theta} &
      a special case of \T{repeat} that only performs rotation between iterations
    \\
    \pSide{n} &
      a side of a regular unit $n$-gon
  \end{tabular}
  \vspace{0.25em}

\noindent
These are expressible in the original DSL,
  and could be even discovered with library learning,
  given an appropriate corpus;
  we treat them as primitives here
  for the sake of simplifying presentation.

\subsection{Representing Equivalent Terms with E-Graphs}

To exploit equivalences during library learning,
  \tname{babble} takes as input a domain-specific equational theory.
For our running example,
  let us assume that the theory contains a single equation:
\begin{equation}
  F \equiv \scale{F}{1} \label{eq:unit-xform}
\end{equation}
  which stipulates that any figure $F$
  is equivalent to itself

  scaled by one.
With this equation in hand,
  it is possible to ``rewrite'' corpus C into corpus A,
  and from there learn the optimal compressed corpus B
  by purely syntactic techniques.
The challenge
  is that there are infinitely many
  alternative corpora C may be rewritten to;
  how do we know which to pick to maximize syntactic alignment,
  and thus the chance of discovering an optimal abstraction?

Instead of trying to guess
  the best equivalent corpus
  or enumerating them one by one,
  \tname{babble} uses \emph{equality saturation}~\cite{tate2009equality,willsey2021egg}.
Equality saturation takes as input
  a term $t$ and a set of rewrite rules,
  and finds the space of all terms equivalent
  to $t$ under the given rules;
  this is possible due to the high degree of
  sharing provided by the \emph{e-graph} data structure,
  which can compactly represent the resulting space.

\autoref{fig:egraph} (left) shows the e-graph
  built by equality saturation for the blue term in \autoref{fig:llmt}---%
  represented in the simplified DSL as \repRot{\pSide{6}}{6}{2\pi/6}---%
  using the rewrite rule \eqref{eq:unit-xform}.
The blue part of the graph represents the original term,
  and the gray part is added by equality saturation.
The solid rectangles in the e-graph denote
  \emph{e-nodes} (which are similar to regular AST nodes),
  while the dashed rectangles denote
  \emph{e-classes} (which represent equivalence classes of terms).
Importantly,
  the edges in the e-graph go from e-nodes to e-classes,
  which enables compact representation of
  programs with variation in sub-terms:
  for example, making e-class $c_1$
  the first child of the \T{repRot} node,
  enables it to represent both terms
    \repRot{\pSide{6}}{6}{2\pi/6}
  and
    \repRot{\scale{\pSide{6}}{1}}{6}{2\pi/6}
  without duplicating their common parts.
Furthermore, because e-graphs can have cycles,
  they can also represent infinite sets of terms:
  for example, the e-class $c_1$ represents
  all terms of the form:
  \pSide{6},
  \scale{\pSide{6}}{1},
  \scale{\scale{\pSide{6}}{1}}{1},
  etc.
Because this e-graph represents
  the space of \emph{all} equivalent terms
  up to the rewrite~\eqref{eq:unit-xform},
  the term we seek for library learning, namely
    \scale{\repRot{\pSide{6}}{6}{2\pi/6}}{1},

  is also represented in the e-class $c_0$.

  \begin{figure}
    \begin{minipage}{.4\textwidth}
      \centering
      \includegraphics[width=.95\textwidth]{figures/egraph-blue.pdf}
    \end{minipage}\vline%
    \begin{minipage}{.6\textwidth}
      \centering
      \includegraphics[width=.95\textwidth]{figures/egraph-new.pdf}
    \end{minipage}
    \caption{(Left) An e-graph representing the space of all programs
    equivalent the blue term in \autoref{fig:llmt} under rule \eqref{eq:unit-xform}.
    %
    (Right) An e-graph with both the blue and the brown terms from \autoref{fig:llmt}
    after equality saturation.
    %
    All terms are written in the simplified DSL.
    }\label{fig:egraph}
  \end{figure}

\subsection{Candidate Generation via E-Graph Anti-Unification}

After building an e-graph from the given corpus by
  running equality saturation with the given equational theory,
  the next step in library learning is
  to generate candidate patterns
  that capture syntactic similarities across the corpus.
The challenge is to generate
  sufficiently few candidates
  to make library learning tractable,
  but sufficiently many to achieve good compression.

We illustrate candidate generation
  using the e-graph in \autoref{fig:egraph} (right),
  which represents the part of our corpus
  consisting of the (saturated) blue and brown terms.
Recall that the optimal
  pattern---which corresponds to the scaled polygon abstraction $f_0$---is:
\begin{equation}
  P_0 = \scale{\repRot{\pSide{X}}{X}{2\pi/X}}{Y}\label{eq:p0}

\end{equation}

Prior work on \tname{DreamCoder} generates patterns by
  picking a random fragment from the corpus,
  and then replacing arbitrarily chosen subterms
  with pattern variables.
For example, to generate the pattern $P_0$,
  \tname{DreamCoder} needs to pick the entire brown subterm as a fragment,
  and then decide to abstract over its subterms $8$ and $2$.
This approach successfully restricts
  the set of candidates from
  all syntactically valid patterns
  to only those that have
  at least one match in the corpus;
  however, since there are too many ways
  to select the subterms to abstract over,
  this space is still too large to explore exhaustively
  beyond small corpora of short programs.

In \tname{babble}, this problem
  is exacerbated by equality saturation,
  since the e-graph often contains exponentially or
  infinitely more programs than the original corpus.

\mypara{Generality Filters}
To further prune the set of viable candidates in \tname{babble},
  we identify two classes of patterns that can be safely discarded,
  either because they are
  too concrete or too abstract.
First,
  a pattern like
    \repRot{\pSide{8}}{X}{2\pi/8}

  can be discarded because it is
  \emph{too concrete} for this corpus:
  the corresponding abstraction
  can be applied only once,
  essentially replacing the sole matching term,
    \repRot{\pSide{8}}{8}{2\pi/8}

  with
    $(\lambda x \ \mathbf{\shortrightarrow}\  \rep{\pSide{8}}{x}{2\pi/8})\ 8$,

  which only adds more AST nodes to the corpus.

Second,
  a pattern like
    \repRot{\pSide{X}}{X}{Y}

  can be discarded because
  it is \emph{too abstract} for this corpus:
  everywhere it matches,
  a more concrete pattern
    \repRot{\pSide{X}}{X}{2\pi/X}

  would also match;
  using the more concrete pattern
  leads to better compression,
  since the actual arguments in its applications
  are both fewer and smaller:
    \T{f 6 2pi/6} and \T{f 8 2pi/8} vs.
    \T{f 6} and \T{f 8}.

Thus,
  our \emphbf{first key insight} is
  to restrict the set of candidates
  to the most concrete patterns that match some pair

  of subterms in the (saturated) corpus.%
  \footnote{As we discuss in \autoref{sec:au} this can in theory eliminate optimal patterns,
  but our evaluation shows that it works well in practice.}
For example, pattern $P_0$
  is the most concrete pattern that matches the two terms
  \begin{gather}
    \scale{\repRot{\pSide{6}}{6}{2\pi/6}}{1}\label{eq:hex}\\
    \scale{\repRot{\pSide{8}}{8}{2\pi/8}}{2}\label{eq:oct}

  \end{gather}
  represented in \autoref{fig:egraph} by the e-classes $c_0$ and $c_2$, respectively.

\mypara{Term Anti-Unification}
Computing the most concrete pattern that
  matches two given terms is known as
  \emph{anti-unification} (AU)~\cite{Reynolds1969TransformationalSA,plotkin1970lattice}.
AU works by
  a simple top-down traversal of the two terms,
  replacing any mismatched constructors by pattern variables.
For example,
  to anti-unify the terms \eqref{eq:hex} and \eqref{eq:oct},
  we start from the root of both terms;
  since both AST nodes share the same constructor \T{scale},
  it becomes part of the pattern and we recurse into the children.
We eventually encounter a mismatch, where
  the term on the left is \T{6} and
  the term on the right is \T{8};
  so we create a fresh pattern variable $X$ and
  remember the \emph{anti-substitution}
  $\sigma = \{(\T{6}, \T{8}) \mapsto X\}$.
When we encounter a mismatch
  in the denominator of the angle,
  we look up the pair of mismatched terms
  $(\T{6}, \T{8})$ in $\sigma$;
  because we already created a variable for this pair,
  we simply return the existing variable $X$.
The final mismatch is $(\T{1}, \T{2})$
  in the second child of \T{scale};
  since this pair is not yet in $\sigma$,
  we create a second pattern variable, $Y$.
At this point,
  the resulting anti-unifier is the pattern $P_0$~\eqref{eq:p0}.

Anti-unifying
  a single pair of terms is simple and efficient.
However,
  in LLMT we want to anti-unify \emph{all pairs of subterms}
  that can occur \emph{in any corpus} equivalent
  (modulo the given equational theory)
  to the original\footnote{
      A careful reader might be wondering if
        we need to compute infinitely many anti-unifiers
        because there might be infinitely many equivalent corpora.
      As we explain shortly,
        there are only finitely many patterns that
        are viable abstraction candidates.
    }.
Explicitly enumerating all equivalent corpora
  represented by the e-graph and then
  performing AU on each pair of subterms
  is infeasible.
Instead, \tname{babble} performs
  anti-unification directly on the e-graph.

\mypara{From Terms to E-Classes}
We first explain how to
  anti-unify two e-classes.
This operation takes as input a pair of e-classes and
  returns a set of \emph{dominant anti-unifiers},
  \emph{i.e.}\xspace a set of patterns that
    (1) match both e-classes, and
    (2) is guaranteed to include
      the best abstraction
      among the most concrete patterns
      that match pairs of terms represented by the two e-classes.

Consider computing $\mathsf{AU}(c_1, c_4)$
  for the e-classes $c_1$ and $c_4$ in
  the e-graph from \autoref{fig:egraph} (right).
AU still proceeds as a top-down traversal,
  but in this context we must
  check whether two e-classes have
  any constructors \emph{in common}.
In this case they do:
  both a \T{side} constructor and a \T{scale} constructor.
Let us first pick the two \T{side} constructors
  and recurse into their only child,
  computing $\mathsf{AU}(\{\T{6}\}, \{\T{8}\})$.
These two e-classes have
  no matching constructors,
  so AU simply returns a pattern variable,
  similarly to term anti-unification;
  this yields the first pattern for $c_1$ and $c_4$: $\T{side}\ X$.

Recall, however, that $c_1$ and $c_4$ also
  have matching \T{scale} constructors.
This is where things get interesting:
  these constructors are involved in a \emph{cycle}
  (their first child is the parent e-class itself).
If we let AU follow this cycle, the set of anti-unifiers becomes infinite:
  $$
    \mathsf{AU}(c_1, c_4) = \{\T{side}\ X\} \cup \{\T{scale}\ p\ 1 \mid p \in \mathsf{AU}(c_1, c_4)\}
  $$
Fortunately,
  we can show that $\T{side}\ X$ \emph{dominates}
  all the other patterns from this set,
  because their pattern variables---in this case, just $X$---match
  the same e-classes, but they are also larger
  (in \autoref{sec:egraphs} we show how
  this domination relation lets us prune other patterns,
  not just those caused by cycles).
Hence $\mathsf{AU}(c_1, c_4)$ simply returns $\{\T{side}\ X\}$.

Following the same logic for the root e-classes of the two polygons, $c_0$ and $c_2$,
  $\mathsf{AU}(c_0, c_2)$ yields that pattern $P_0$~\eqref{eq:p0},
  which is required to learn the optimal abstraction.

\mypara{From E-Classes to E-Graphs}
To obtain the set of all candidate abstractions,
we need to perform anti-unification over all pairs of e-classes in the e-graph.
Clearly, these computations have overlapping subproblems
(for example, we have to compute $\mathsf{AU}(c_1, c_4)$
as part of $\mathsf{AU}(c_0, c_2)$ and  $\mathsf{AU}(c_0, c_3)$).
To avoid duplicating work, \tname{babble} uses an efficient dynamic programming algorithm
that processes all pairs of e-classes in a bottom-up fashion.

\subsection{Candidate Selection via Targeted Common Subexpression Elimination}

We now illustrate the final step of library learning in \tname{babble}:
given the set of candidate abstractions generated by e-graph anti-unification,
the goal is to pick a subset that gives the best compression for the corpus as a whole.
For example, the candidates generated for the corpus from \autoref{fig:llmt} include:
\[ \begin{array}{rcll}
  f_0 &=& \lambda x\ y \ \mathbf{\shortrightarrow}\  \scale{\repRot{\pSide{x}}{x}{2\pi/x}}{y} &\quad \text{scaled regular $x$-gon} \\[2pt]
  f_1 &=& \lambda x \ \mathbf{\shortrightarrow}\  \repRot{\pSide{x}}{x}{2\pi/x} &\quad \text{regular $x$-gon} \\[2pt]
  f_2 &=& \lambda y \ \mathbf{\shortrightarrow}\  \scale{\repRot{\pSide{6}}{6}{2\pi/6}}{y} &\quad \text{scaled hexagon}
\end{array} \]
It is not immediately clear which abstraction is better:
  $f_0$ matches more terms than $f_2$,
  but $f_2$ requires fewer arguments
  (so if we have enough scaled hexagons in the corpus and
  only one octagon, it might be better to leave the
  octagon un-abstracted).
On the other hand,
  $f_1$ might be better,
  since it does not introduce the
  redundant transformation on the blue hexagons.
Finally,
  if we pick $f_0$ \emph{and} $f_2$ together,
  we can also abstract the definition of
  $f_2$ as $f_0\ 6$, thereby getting additional reuse!
As you can see,
  candidate selection is highly non-trivial,
  since it needs to take into account the
  choice between different equivalent programs in the e-graph,
  and the fact that some abstractions can be defined using others.

\begin{figure}
  \centering
  \includegraphics[width=.8\textwidth]{figures/beam.pdf}
  \caption{The e-graph from \autoref{fig:egraph} (right) with applications of $f_0$ and $f_1$ depicted in red.
           We show the unchanged parts of the graph representing the unit hexagon and octagon as corresponding shapes;
           we also omit some of the gray e-nodes added in the previous stage.}\label{fig:beam}
\end{figure}

\mypara{Reduction to E-Graph Extraction}
To overcome this difficulty, we once again leverage e-graph and equality saturation.
Our \emphbf{second key insight} is that selecting the optimal subset of abstractions
can be reduced to the problem of extracting the smallest term from an e-graph
in the presence of common sub-expressions.

To illustrate this reduction,
let us limit our attention to only two candidate abstractions, $f_0$ and $f_1$, defined above.
\tname{babble} converts each of the candidate patterns and its corresponding abstraction into a rewrite rule
that introduces a \emph{local} $\lambda$-abstraction followed by application into the corpus;
for our two abstractions these rules are:
\begin{align}

  \scale{\repRot{\pSide{X}}{X}{2\pi/X}}{Y} &\Rightarrow f_0\ X\ Y\label{eq:rw1}\\
  \repRot{\pSide{X}}{X}{2\pi/X} &\Rightarrow f_1\ X \label{eq:rw2}
\end{align}
The result of applying these rules to the e-graph from \autoref{fig:egraph} (right)
is shown in \autoref{fig:beam}.
For example, you can see that the e-class $c_2$ (which represents the scale-2 octagon)
now stores an alternative representation:
$f_0\  8\ 2$.
The e-class $c_0$ (unit hexagon)
has representations using either $f_0$ or $f_1$,
because this class matches both rewrite rules \eqref{eq:rw1} and \eqref{eq:rw2} above.
Note also that because the definitions of the $\lambda$-abstractions are also stored in the e-graph,
equality saturation can use the above rewrite rules inside their bodies,
to make one function use another:
for example, one term stored for the definition of $f_0$ is
$$
\lambda x\ y \ \mathbf{\shortrightarrow}\  \scale{f_1\ x}{y}.
$$

Once we have built the version of the e-graph with local lambdas for all the candidate abstractions,
all that is left is to extract the smallest term out of this e-graph.
The tricky part is that we want to count the size of the duplicated lambdas only once.
For example, in \autoref{fig:beam}, $f_0$ is applied twice (in $c_0$ and $c_2$);
if term extraction were to choose both of these e-nodes,
we want to treat their first child (the definition of $f_0$) as a common sub-expression,
whose size contributes to the final expression only once.
Intuitively, this is because we can ``float'' these lambdas
into top-level \T{let}-bindings,
thereby defining $f_0$ only once, and replacing each local lambda with a name.

Extraction with common sub-expressions is a known but notoriously hard problem,
which is traditionally reduced to integer linear programming (ILP)~\cite{tensat,spores}.
Because the scalability of the ILP approach is very limited,
we have developed a custom extraction algorithm,
which scales better by using domain-specific knowledge and approximation.

\mypara{Extraction Algorithm}
The main idea for making extraction more efficient is that
for library learning we are only interested in sharing a certain type of sub-terms:
namely, the $\lambda$-abstractions.
Hence for each e-class we only need to keep track
of the the best term for each possible library
(\emph{i.e.}\xspace each subset of $\lambda$-abstractions).
More precisely, for each e-class and library,
we keep track of
  (1) the smallest size of the library
  (2) the smallest size of the program refactored using this library
  (counting the $\lambda$-abstractions as a single node).
We compute and propagate this information bottom-up through the e-graph.
Once this information is computed for the root e-class
(that represents the entire corpus),
we can simply choose the library with the smallest total size.

For example, for the e-class $c_2$ from \autoref{fig:beam},
with the empty library $\emptyset$, the size of library is 0 and the size of the smallest program (\scale{\repRot{\pSide{8}}{8}{2\pi/8}}{2}) is 9;
with the library $\{f_0\}$, the size of the library is 9 and the size of the smallest program ($f_0\ 2\ 8$) is 3;
while with the library $\{f_1\}$, the size of the library is 7 and the size of the smallest program (\scale{f_1\ 8}{2}) is 4.
Clearly for this e-class introducing library functions is not paying off yet ($0 + 9 < 7 + 4 < 9 + 3$),
since each one can be only used once.
This situation changes, however, as we move up towards the root.
Already for the parent e-class of $c_0$ and $c_2$,
the cost of introducing $\{f_0\}$ and $\{f_1\}$ is amortized:
the size of the smallest program is 17 with the empty library and 7 with either $\{f_0\}$ or $\{f_1\}$,
so $\{f_1\}$ is already worthwhile to introduce ($9 + 7 < 0 + 17$).
Including even more programs with scaled polygons eventually makes the library $\{f_0\}$ the most profitable of the four subsets.

Since in a larger corpus, keeping track of all subsets of candidate abstractions is not feasible,
\tname{babble} provides a way to trade off scalability and precision
by using a \emph{beam search} approach.

\section{Related Work}
\label{sec:related}

\tname{babble} is inspired by work on library learning,
specifically the \tname{DreamCoder} line of work,
as well as equality saturation-based program synthesis and decompilation.

\mypara{\tname{DreamCoder}}
\tname{DreamCoder}~\cite{dreamcoder} is a program synthesizer
  that learns a library of abstractions
  from solutions to a set of related synthesis tasks.
The library is intended to be used for solving other
  similar synthesis tasks.
\tname{DreamCoder}
  uses version spaces~\cite{vsa1, vsa2}
  to compactly store a large
  set of programs
  and leverages ideas from e-graphs\xspace
  (such as e-matching)
  but only for exploring the space of refactorings
  of the original program using the candidate libraries,
  not for making library learning robust to syntactic variation.
Our evaluation shows that \tname{babble} can find more optimal abstractions faster than \tname{DreamCoder}.

\tname{DreamCoder} has sparked several direction of follow-up work
that attempt to improve the efficiency of its library learning procedure
and the quality of the learned abstractions.
One of them is by \citet{laps},
which uses natural language annotations and a neural network
to guide library learning.
Another one is \tname{Stitch}~\cite{stitch},
which was developed \emph{concurrently} with this work
and is the most closely related to \tname{babble};
we discuss \tname{Stitch} in some detail below.

\mypara{\tname{Stitch}}
The core difference between the two approaches
is that \tname{Stitch} focuses on improving \emph{efficiency} of purely syntactic library learning,
whereas \tname{babble} attempts to improve its \emph{expressiveness} by adding equational theories.
While \tname{babble} separates library learning into two phases---%
candidate generation via anti-unification and candidate selection via e-graph extraction---%
the \tname{Stitch} algorithm \emph{interleaves} the generation and selection phases
in a branch-and-bound top-down search.
Starting from the ``top'' pattern $X$,
\tname{Stitch} gradually refines it until further refinement does not pay off.
To quickly prune suboptimal candidate patterns,
\tname{Stitch} computes an upper bound on their compression
by summing up the compression at each match of this pattern in the corpus
(this bound is imprecise because it does not take into account that matches might overlap).
For candidates that are not pruned this way,
\tname{Stitch} computes their true compression by searching for the optimal subset of matches to rewrite,
a so-called ``rewrite strategy''.
\tname{babble}'s extraction algorithm can be seen as a generalization of \tname{Stitch}'s rewrite strategy:
while the former searches over both subsets of patterns and how to apply them to the corpus at the same time,
the latter considers a single pattern at a time and only searches for the best way to apply it.
Since the search space in the former case is much larger,
\tname{babble} uses a beam search approximation,
while in \tname{Stitch} the rewrite strategy is precise.
To sum up, the main pros and cons of the two approaches are:
\begin{itemize}
\item \tname{babble} can learn libraries modulo equational theories, while \tname{Stitch} cannot;
\item \tname{Stitch} provides optimality guarantees for learning a single best abstraction at a time, while \tname{babble} can learn multiple abstractions at once, but sacrifices theoretical optimality.
\end{itemize}

\mypara{Other Library Learning Techniques}
\tname{Knorf}~\cite{knorf} is a library learning tool for logic programs,
which, like \tname{babble}, proceeds in two phases.
Their candidate generation phase is similar to the upper bound computation in \tname{Stitch},
while their selection phase uses an off-the-shelf constraint solver.
It would be interesting to explore whether their constraint-based technique can be generalized beyond logic programs.

Other work develops limited forms of library learning,
where only certain kinds of sub-terms can be abstracted.
For example, ShapeMod~\cite{shapemod} learns macros
  for 3D shapes represented
  in a DSL called ShapeAssembly,
  and only supports abstracting over numeric parameters,
  like dimensions of shapes.
Our own prior work~\cite{cogsci-smiley} extracts common structure from graphical programs,
  but only supports abstracting over primitive shapes
  and applying the abstraction at the top level of the program.
Such restrictions make the library learning problem more computationally tractable,
  but limit the expressiveness of the learned abstractions.

There are several neural program synthesis tools%
  ~\cite{polozov, srini, learner, gredilla}
  that learn \emph{programming idioms}
  using statistical techniques.

Some of these tools have used
  ``explore-compress'' algorithms~\cite{learner} to
  iteratively enumerate a set of
  programs from a grammar and find a solution
  that exposes abstractions that make the set of
  programs maximally compressible.
This is similar to
  common subexpression elimination which \tname{babble} uses
  for guiding extraction.

\mypara{Loop rerolling}
Loop rerolling is related to library learning in that it also aims to discover hidden structure in a program,
except that this structure is in the form of loops.
A variety of domains have used loop rerolling
  to infer abstractions from flat input programs.
In hardware, loop rerolling is
  used to optimize programs
  for code size~\cite{hardware1, hardware2, hardware3}.
In many of these tools,
  the compiler first unrolls a loop,
  applies optimizations, then rerolls it --- the compiler
  therefore has structural information about the loop
  that can be used for rerolling~\cite{hardware1}.
The graphics domain has used
  loop-rerolling to discover
  latent structure from low-level representations.
CSGNet~\cite{csgnet} and Shape2Prog~\cite{shape2prog} used
  neural program generators to discover for loops from
  pixel- and voxel-based input representations.
\cite{handdrawn} used program synthesis
  and machine learning
  to infer loops from hand-drawn images.
Szalinski~\cite{szalinski} used
  equality saturation to
  automatically learn loops in the form
  of maps and folds
  from flat 3D CAD programs that are
  synthesized by mesh decompilation tools~\cite{reincarnate}.
WebRobot~\cite{webrobot} has used speculative
  rewriting for inferring loops from traces
  of web interactions.
Similar to \tname{babble} (and unlike Szalinski),
  WebRobot finds abstractions
  over \textit{multiple} input traces.

\mypara{Applications of Anti-Unification}
Anti-unification is a well-established technique for discovering common structure in programs.
It is the core idea behind bottom-up Inductive Logic Programming~\cite{ilp},
and has also been used for software clone detection~\cite{Bulychev10},
programming by example~\cite{raza2014programming},
and learning repetitive code edits~\cite{lase,refaser}.
It is possible that these applications could also benefit from \tname{babble}'s notion of anti-unification over e-graphs
to make them more robust to semantics-preserving transformations.

\mypara{Synthesis and Optimization using E-graphs\xspace}
While traditionally e-graphs\xspace have been used
  in SMT solvers for
  facilatiting communication between
  different theories,
  several tools have demonstrated
  their use for optimization
  and synthesis.
\citet{tate2009equality} first used e-graphs\xspace for
  equality saturation: a rewrite-driven technique
  for optimizing Java programs with loops.
Since then, several tools have used equality saturation
  for finding programs equivalent to,
  but better than, some
  input program~\cite{willsey2021egg, herbie, tensat, szalinski, diospyros, spores,
  cc1}.
\tname{babble} uses an anti-unification
  algorithm on e-graphs\xspace
  (together with domain specific
  rewrites), which prior work has not shown.
Additionally, prior work has either used
  greedy or ILP-based extraction strategies,
  whereas \tname{babble} uses a new targeted common
  subexpression elimination approach which
  we believe can be used in many other applications of
  equality saturation, especially given its
  amenability to approximation via beam search.