Add lambda lift utility (#7601)

Towards a lazy approach for handling lambdas in the higher-order extension.

6 files changed, 328 insertions, 0 deletions

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index c145f7fda..db18a4bc8 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -1121,6 +1121,8 @@ libcvc5_add_sources(
   theory/uf/eq_proof.h
   theory/uf/function_const.cpp
   theory/uf/function_const.h
+  theory/uf/lambda_lift.cpp
+  theory/uf/lambda_lift.h
   theory/uf/proof_checker.cpp
   theory/uf/proof_checker.h
   theory/uf/proof_equality_engine.cpp
diff --git a/src/options/uf_options.toml b/src/options/uf_options.toml
index 7392dd50c..a289132fc 100644
--- a/src/options/uf_options.toml
+++ b/src/options/uf_options.toml
@@ -59,3 +59,11 @@ name   = "Uninterpreted Functions Theory"
   type       = "bool"
   default    = "true"
   help       = "apply extensionality on function symbols"
+
+[[option]]
+  name       = "ufHoLazyLambdaLift"
+  category   = "regular"
+  long       = "uf-lazy-ll"
+  type       = "bool"
+  default    = "false"
+  help       = "do lambda lifting lazily"
diff --git a/src/proof/eager_proof_generator.cpp b/src/proof/eager_proof_generator.cpp
index 2a86f982c..36ef1bf58 100644
--- a/src/proof/eager_proof_generator.cpp
+++ b/src/proof/eager_proof_generator.cpp
@@ -134,6 +134,18 @@ TrustNode EagerProofGenerator::mkTrustedRewrite(Node a,
   return TrustNode::mkTrustRewrite(a, b, this);
 }
 
+TrustNode EagerProofGenerator::mkTrustedRewrite(Node a,
+                                                Node b,
+                                                PfRule id,
+                                                const std::vector<Node>& args)
+{
+  Node eq = a.eqNode(b);
+  CDProof cdp(d_pnm);
+  cdp.addStep(eq, id, {}, args);
+  std::shared_ptr<ProofNode> pf = cdp.getProofFor(eq);
+  return mkTrustedRewrite(a, b, pf);
+}
+
 TrustNode EagerProofGenerator::mkTrustedPropagation(
     Node n, Node exp, std::shared_ptr<ProofNode> pf)
 {
diff --git a/src/proof/eager_proof_generator.h b/src/proof/eager_proof_generator.h
index 90e2787fd..9ebc89bd0 100644
--- a/src/proof/eager_proof_generator.h
+++ b/src/proof/eager_proof_generator.h
@@ -158,6 +158,21 @@ class EagerProofGenerator : public ProofGenerator
    * a proof of a = b
    */
   TrustNode mkTrustedRewrite(Node a, Node b, std::shared_ptr<ProofNode> pf);
+  /**
+   * Make trust node from a single step proof. This is a convenience function
+   * that avoids the need to explictly construct ProofNode by the caller.
+   *
+   * @param a the original
+   * @param b what is rewrites to
+   * @param id The rule of the proof concluding a=b
+   * @param args The arguments to the proof concluding a=b,
+   * @return The trust node corresponding to the fact that this generator has
+   * a proof of a=b.
+   */
+  TrustNode mkTrustedRewrite(Node a,
+                             Node b,
+                             PfRule id,
+                             const std::vector<Node>& args);
   //--------------------------------------- common proofs
   /**
    * This returns the trust node corresponding to the splitting lemma
diff --git a/src/theory/uf/lambda_lift.cpp b/src/theory/uf/lambda_lift.cpp
new file mode 100644
index 000000000..3d51e6858
--- /dev/null
+++ b/src/theory/uf/lambda_lift.cpp
@@ -0,0 +1,194 @@
+/******************************************************************************
+ * Top contributors (to current version):
+ *   Andrew Reynolds
+ *
+ * This file is part of the cvc5 project.
+ *
+ * Copyright (c) 2009-2021 by the authors listed in the file AUTHORS
+ * in the top-level source directory and their institutional affiliations.
+ * All rights reserved.  See the file COPYING in the top-level source
+ * directory for licensing information.
+ * ****************************************************************************
+ *
+ * Implementation of lambda lifting.
+ */
+
+#include "theory/uf/lambda_lift.h"
+
+#include "expr/node_algorithm.h"
+#include "expr/skolem_manager.h"
+#include "options/uf_options.h"
+#include "smt/env.h"
+
+using namespace cvc5::kind;
+
+namespace cvc5 {
+namespace theory {
+namespace uf {
+
+LambdaLift::LambdaLift(Env& env)
+    : EnvObj(env),
+      d_lifted(userContext()),
+      d_lambdaMap(userContext()),
+      d_epg(env.isTheoryProofProducing() ? new EagerProofGenerator(
+                env.getProofNodeManager(), userContext(), "LambdaLift::epg")
+                                         : nullptr)
+{
+}
+
+TrustNode LambdaLift::lift(Node node)
+{
+  if (d_lifted.find(node) != d_lifted.end())
+  {
+    return TrustNode::null();
+  }
+  d_lifted.insert(node);
+  Node assertion = getAssertionFor(node);
+  if (assertion.isNull())
+  {
+    return TrustNode::null();
+  }
+  // if no proofs, return lemma with no generator
+  if (d_epg == nullptr)
+  {
+    return TrustNode::mkTrustLemma(assertion);
+  }
+  return d_epg->mkTrustNode(
+      assertion, PfRule::MACRO_SR_PRED_INTRO, {}, {assertion});
+}
+
+TrustNode LambdaLift::ppRewrite(Node node, std::vector<SkolemLemma>& lems)
+{
+  TNode skolem = getSkolemFor(node);
+  if (skolem.isNull())
+  {
+    return TrustNode::null();
+  }
+  d_lambdaMap[skolem] = node;
+  if (!options().uf.ufHoLazyLambdaLift)
+  {
+    TrustNode trn = lift(node);
+    lems.push_back(SkolemLemma(trn, skolem));
+  }
+  // if no proofs, return lemma with no generator
+  if (d_epg == nullptr)
+  {
+    return TrustNode::mkTrustRewrite(node, skolem);
+  }
+  Node eq = node.eqNode(skolem);
+  return d_epg->mkTrustedRewrite(
+      node, skolem, PfRule::MACRO_SR_PRED_INTRO, {eq});
+}
+
+Node LambdaLift::getLambdaFor(TNode skolem) const
+{
+  NodeNodeMap::const_iterator it = d_lambdaMap.find(skolem);
+  if (it == d_lambdaMap.end())
+  {
+    return Node::null();
+  }
+  return it->second;
+}
+
+Node LambdaLift::getAssertionFor(TNode node)
+{
+  TNode skolem = getSkolemFor(node);
+  if (skolem.isNull())
+  {
+    return Node::null();
+  }
+  Kind k = node.getKind();
+  Node assertion;
+  if (k == LAMBDA)
+  {
+    NodeManager* nm = NodeManager::currentNM();
+    // The new assertion
+    std::vector<Node> children;
+    // bound variable list
+    children.push_back(node[0]);
+    // body
+    std::vector<Node> skolem_app_c;
+    skolem_app_c.push_back(skolem);
+    skolem_app_c.insert(skolem_app_c.end(), node[0].begin(), node[0].end());
+    Node skolem_app = nm->mkNode(APPLY_UF, skolem_app_c);
+    children.push_back(skolem_app.eqNode(node[1]));
+    // axiom defining skolem
+    assertion = nm->mkNode(FORALL, children);
+
+    // Lambda lifting is trivial to justify, hence we don't set a proof
+    // generator here. In particular, replacing the skolem introduced
+    // here with its original lambda ensures the new assertion rewrites
+    // to true.
+    // For example, if (lambda y. t[y]) has skolem k, then this lemma is:
+    //   forall x. k(x)=t[x]
+    // whose witness form rewrites
+    //   forall x. (lambda y. t[y])(x)=t[x] --> forall x. t[x]=t[x] --> true
+  }
+  return assertion;
+}
+
+Node LambdaLift::getSkolemFor(TNode node)
+{
+  Node skolem;
+  Kind k = node.getKind();
+  if (k == LAMBDA)
+  {
+    // if a lambda, return the purification variable for the node. We ignore
+    // lambdas with free variables, which can occur beneath quantifiers
+    // during preprocessing.
+    if (!expr::hasFreeVar(node))
+    {
+      Trace("rtf-proof-debug")
+          << "RemoveTermFormulas::run: make LAMBDA skolem" << std::endl;
+      // Make the skolem to represent the lambda
+      NodeManager* nm = NodeManager::currentNM();
+      SkolemManager* sm = nm->getSkolemManager();
+      skolem = sm->mkPurifySkolem(
+          node,
+          "lambdaF",
+          "a function introduced due to term-level lambda removal");
+    }
+  }
+  return skolem;
+}
+
+TrustNode LambdaLift::betaReduce(TNode node) const
+{
+  Kind k = node.getKind();
+  if (k == APPLY_UF)
+  {
+    Node op = node.getOperator();
+    Node opl = getLambdaFor(op);
+    if (!opl.isNull())
+    {
+      std::vector<Node> args(node.begin(), node.end());
+      Node app = betaReduce(opl, args);
+      Trace("uf-lazy-ll") << "Beta reduce: " << node << " -> " << app
+                          << std::endl;
+      if (d_epg == nullptr)
+      {
+        return TrustNode::mkTrustRewrite(node, app);
+      }
+      return d_epg->mkTrustedRewrite(
+          node, app, PfRule::MACRO_SR_PRED_INTRO, {node.eqNode(app)});
+    }
+  }
+  // otherwise, unchanged
+  return TrustNode::null();
+}
+
+Node LambdaLift::betaReduce(TNode lam, const std::vector<Node>& args) const
+{
+  Assert(lam.getKind() == LAMBDA);
+  NodeManager* nm = NodeManager::currentNM();
+  std::vector<Node> betaRed;
+  betaRed.push_back(lam);
+  betaRed.insert(betaRed.end(), args.begin(), args.end());
+  Node app = nm->mkNode(APPLY_UF, betaRed);
+  app = rewrite(app);
+  return app;
+}
+
+}  // namespace uf
+}  // namespace theory
+}  // namespace cvc5
diff --git a/src/theory/uf/lambda_lift.h b/src/theory/uf/lambda_lift.h
new file mode 100644
index 000000000..f61007d1c
--- /dev/null
+++ b/src/theory/uf/lambda_lift.h
@@ -0,0 +1,97 @@
+/******************************************************************************
+ * Top contributors (to current version):
+ *   Andrew Reynolds
+ *
+ * This file is part of the cvc5 project.
+ *
+ * Copyright (c) 2009-2021 by the authors listed in the file AUTHORS
+ * in the top-level source directory and their institutional affiliations.
+ * All rights reserved.  See the file COPYING in the top-level source
+ * directory for licensing information.
+ * ****************************************************************************
+ *
+ * Lambda lifting
+ */
+
+#include "cvc5_private.h"
+
+#ifndef CVC5__THEORY__UF__LAMBDA_LIFT_H
+#define CVC5__THEORY__UF__LAMBDA_LIFT_H
+
+#include "context/cdhashmap.h"
+#include "context/cdhashset.h"
+#include "expr/node.h"
+#include "proof/eager_proof_generator.h"
+#include "proof/trust_node.h"
+#include "smt/env_obj.h"
+#include "theory/skolem_lemma.h"
+
+namespace cvc5 {
+namespace theory {
+namespace uf {
+
+/**
+ * Module for doing various operations on lambdas, including lambda lifting.
+ *
+ * In the following, we say a "lambda function" is a skolem variable that
+ * was introduced as a purification skolem for a lambda term.
+ */
+class LambdaLift : protected EnvObj
+{
+  typedef context::CDHashSet<Node> NodeSet;
+  typedef context::CDHashMap<Node, Node> NodeNodeMap;
+
+ public:
+  LambdaLift(Env& env);
+
+  /**
+   * process, return the trust node corresponding to the lemma for the lambda
+   * lifting of (lambda) term node, or null if it is not a lambda or if
+   * the lambda lifting lemma has already been generated in this context.
+   */
+  TrustNode lift(Node node);
+
+  /**
+   * This method has the same contract as Theory::ppRewrite.
+   * Preprocess, return the trust node corresponding to the rewrite. A null
+   * trust node indicates no rewrite.
+   */
+  TrustNode ppRewrite(Node node, std::vector<SkolemLemma>& lems);
+
+  /** Get the lambda term for skolem, if skolem is a lambda function. */
+  Node getLambdaFor(TNode skolem) const;
+
+  /**
+   * Beta-reduce node. If node is APPLY_UF and its operator is a lambda
+   * function known by this class, then this method returns the beta
+   * reduced version of node. We only beta-reduce the top-most application
+   * in node.
+   *
+   * This method returns the trust node corresponding to the rewrite of node to
+   * the return value. It returns the null trust node if no beta reduction is
+   * possible for node.
+   */
+  TrustNode betaReduce(TNode node) const;
+  /** Beta-reduce the given lambda on the given arguments. */
+  Node betaReduce(TNode lam, const std::vector<Node>& args) const;
+
+ private:
+  /**
+   * Get assertion for node, which is the axiom defining
+   */
+  static Node getAssertionFor(TNode node);
+  /** Get skolem for lambda term node, returns its purification skolem */
+  static Node getSkolemFor(TNode node);
+  /** The nodes we have already returned trust nodes for */
+  NodeSet d_lifted;
+  /** Mapping skolems to their lambda */
+  NodeNodeMap d_lambdaMap;
+  /** An eager proof generator */
+  std::unique_ptr<EagerProofGenerator> d_epg;
+};
+
+}  // namespace uf
+}  // namespace theory
+}  // namespace cvc5
+
+#endif /* CVC5__THEORY__UF__LAMBDA_LIFT_H */