Add lazy approach for handling lambdas in the HO extension (#7625)

This adds a new option --uf-lazy-ll for not applying lambda lifting eagerly. This is motivated by HO operators in theory of bags, e.g. bag.map, which cannot answer "sat" for simple problems where lambdas are mapped to bags, due to the introduction of quantified formulas for eliminating lambdas.

It moves lambda lifting from term formula removal to a utility module within TheoryUF. If lazy lambda lifting is enabled, this module does not introduce quantified formulas for lambdas eagerly.

It adds 2 lemma schemas for reasoning about quantifier-free constraints with lambdas natively in TheoryUF. It extends the model construction in the HoExtension to assign lambdas in the case that an equivalence class contains a variable that is defined by a lambda.

19 files changed, 439 insertions, 88 deletions

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index db18a4bc8..4d90ae742 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -1129,6 +1129,8 @@ libcvc5_add_sources(
   theory/uf/proof_equality_engine.h
   theory/uf/ho_extension.cpp
   theory/uf/ho_extension.h
+  theory/uf/lambda_lift.cpp
+  theory/uf/lambda_lift.h
   theory/uf/symmetry_breaker.cpp
   theory/uf/symmetry_breaker.h
   theory/uf/theory_uf.cpp
diff --git a/src/api/cpp/cvc5_kind.h b/src/api/cpp/cvc5_kind.h
index 90a57317d..c0c269be2 100644
--- a/src/api/cpp/cvc5_kind.h
+++ b/src/api/cpp/cvc5_kind.h
@@ -331,10 +331,6 @@ enum Kind : int32_t
    *   - `Solver::mkTerm(Kind kind, const std::vector<Term>& children) const`
    */
   CARDINALITY_CONSTRAINT,
-#if 0
-  /* Partial uninterpreted function application.  */
-  PARTIAL_APPLY_UF,
-#endif
   /**
    * Higher-order applicative encoding of function application, left
    * associative.
diff --git a/src/smt/term_formula_removal.cpp b/src/smt/term_formula_removal.cpp
index edf4e2761..341893a70 100644
--- a/src/smt/term_formula_removal.cpp
+++ b/src/smt/term_formula_removal.cpp
@@ -342,48 +342,6 @@ Node RemoveTermFormulas::runCurrentInternal(TNode node,
       }
     }
   }
-  else if (node.getKind() == kind::LAMBDA)
-  {
-    // if a lambda, do lambda-lifting
-    if (!expr::hasFreeVar(node))
-    {
-      skolem = getSkolemForNode(node);
-      if (skolem.isNull())
-      {
-        Trace("rtf-proof-debug")
-            << "RemoveTermFormulas::run: make LAMBDA skolem" << std::endl;
-        // Make the skolem to represent the lambda
-        SkolemManager* sm = nodeManager->getSkolemManager();
-        skolem = sm->mkPurifySkolem(
-            node,
-            "lambdaF",
-            "a function introduced due to term-level lambda removal");
-        d_skolem_cache.insert(node, skolem);
-
-        // 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 = nodeManager->mkNode(kind::APPLY_UF, skolem_app_c);
-        children.push_back(skolem_app.eqNode(node[1]));
-        // axiom defining skolem
-        newAssertion = nodeManager->mkNode(kind::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
-      }
-    }
-  }
   else if (node.getKind() == kind::WITNESS)
   {
     // If a witness choice
diff --git a/src/theory/inference_id.cpp b/src/theory/inference_id.cpp
index 653762721..688014574 100644
--- a/src/theory/inference_id.cpp
+++ b/src/theory/inference_id.cpp
@@ -451,6 +451,8 @@ const char* toString(InferenceId i)
     case InferenceId::UF_HO_MODEL_APP_ENCODE: return "UF_HO_MODEL_APP_ENCODE";
     case InferenceId::UF_HO_MODEL_EXTENSIONALITY:
       return "UF_HO_MODEL_EXTENSIONALITY";
+    case InferenceId::UF_HO_LAMBDA_UNIV_EQ: return "HO_LAMBDA_UNIV_EQ";
+    case InferenceId::UF_HO_LAMBDA_APP_REDUCE: return "HO_LAMBDA_APP_REDUCE";
 
     default: return "?";
   }
diff --git a/src/theory/inference_id.h b/src/theory/inference_id.h
index da2805df0..a13bc4d1b 100644
--- a/src/theory/inference_id.h
+++ b/src/theory/inference_id.h
@@ -838,6 +838,15 @@ enum class InferenceId
   // different applications
   //   (not (= (f sk1 .. skn) (g sk1 .. skn))
   UF_HO_MODEL_EXTENSIONALITY,
+  // equivalence of lambda functions
+  //   f = g => forall x. reduce(lambda(f)(x)) = reduce(lambda(g)(x))
+  // This is applied when lamda functions f and g are in the same eq class.
+  UF_HO_LAMBDA_UNIV_EQ,
+  // equivalence of a lambda function and an ordinary function
+  //   f = h => h(t) = reduce(lambda(f)(t))
+  // This is applied when lamda function f and ordinary function h are in the
+  // same eq class.
+  UF_HO_LAMBDA_APP_REDUCE,
   //-------------------- end model-construction specific part
   //-------------------- end HO extension to UF
   //-------------------------------------- end uf theory
diff --git a/src/theory/uf/ho_extension.cpp b/src/theory/uf/ho_extension.cpp
index 96029eab8..aec223d66 100644
--- a/src/theory/uf/ho_extension.cpp
+++ b/src/theory/uf/ho_extension.cpp
@@ -19,6 +19,7 @@
 #include "expr/skolem_manager.h"
 #include "options/uf_options.h"
 #include "theory/theory_model.h"
+#include "theory/uf/lambda_lift.h"
 #include "theory/uf/theory_uf_rewriter.h"
 
 using namespace std;
@@ -30,31 +31,83 @@ namespace uf {
 
 HoExtension::HoExtension(Env& env,
                          TheoryState& state,
-                         TheoryInferenceManager& im)
+                         TheoryInferenceManager& im,
+                         LambdaLift& ll)
     : EnvObj(env),
       d_state(state),
       d_im(im),
+      d_ll(ll),
       d_extensionality(userContext()),
+      d_cachedLemmas(userContext()),
       d_uf_std_skolem(userContext())
 {
   d_true = NodeManager::currentNM()->mkConst(true);
 }
 
-Node HoExtension::ppRewrite(Node node)
+TrustNode HoExtension::ppRewrite(Node node, std::vector<SkolemLemma>& lems)
 {
-  // convert HO_APPLY to APPLY_UF if fully applied
-  if (node.getKind() == HO_APPLY)
+  Kind k = node.getKind();
+  if (k == HO_APPLY)
   {
+    // convert HO_APPLY to APPLY_UF if fully applied
     if (node[0].getType().getNumChildren() == 2)
     {
       Trace("uf-ho") << "uf-ho : expanding definition : " << node << std::endl;
       Node ret = getApplyUfForHoApply(node);
       Trace("uf-ho") << "uf-ho : ppRewrite : " << node << " to " << ret
                      << std::endl;
-      return ret;
+      return TrustNode::mkTrustRewrite(node, ret);
     }
+    // partial beta reduction
+    // f ---> (lambda ((x Int) (y Int)) s[x, y]) then (@ f t) is preprocessed
+    // to (lambda ((y Int)) s[t, y]).
+    if (options().uf.ufHoLazyLambdaLift)
+    {
+      Node op = node[0];
+      Node opl = d_ll.getLambdaFor(op);
+      if (!opl.isNull())
+      {
+        NodeManager* nm = NodeManager::currentNM();
+        Node app = nm->mkNode(HO_APPLY, opl, node[1]);
+        app = rewrite(app);
+        Trace("uf-lazy-ll")
+            << "Partial beta reduce: " << node << " -> " << app << std::endl;
+        return TrustNode::mkTrustRewrite(node, app, nullptr);
+      }
+    }
+  }
+  else if (k == APPLY_UF)
+  {
+    // Say (lambda ((x Int)) t[x]) occurs in the input. We replace this
+    // by k during ppRewrite. In the following, if we see (k s), we replace
+    // it by t[s]. This maintains the invariant that the *only* occurences
+    // of k are as arguments to other functions; k is not applied
+    // in any preprocessed constraints.
+    if (options().uf.ufHoLazyLambdaLift)
+    {
+      // if an application of the lambda lifted function, do beta reduction
+      // immediately
+      Node op = node.getOperator();
+      Node opl = d_ll.getLambdaFor(op);
+      if (!opl.isNull())
+      {
+        Assert(opl.getKind() == LAMBDA);
+        std::vector<Node> args(node.begin(), node.end());
+        Node app = d_ll.betaReduce(opl, args);
+        Trace("uf-lazy-ll")
+            << "Beta reduce: " << node << " -> " << app << std::endl;
+        return TrustNode::mkTrustRewrite(node, app, nullptr);
+      }
+    }
+  }
+  else if (k == kind::LAMBDA)
+  {
+    Trace("uf-lazy-ll") << "Preprocess lambda: " << node << std::endl;
+    TrustNode skTrn = d_ll.ppRewrite(node, lems);
+    Trace("uf-lazy-ll") << "...return " << skTrn.getNode() << std::endl;
+    return skTrn;
   }
-  return node;
+  return TrustNode::null();
 }
 
 Node HoExtension::getExtensionalityDeq(TNode deq, bool isCached)
@@ -217,7 +270,7 @@ unsigned HoExtension::checkExtensionality(TheoryModel* m)
   {
     Node eqc = (*eqcs_i);
     TypeNode tn = eqc.getType();
-    if (tn.isFunction())
+    if (tn.isFunction() && d_lambdaEqc.find(eqc) == d_lambdaEqc.end())
     {
       hasFunctions = true;
       // if during collect model, must have an infinite type
@@ -413,6 +466,168 @@ unsigned HoExtension::checkAppCompletion()
   return 0;
 }
 
+unsigned HoExtension::checkLazyLambda()
+{
+  if (!options().uf.ufHoLazyLambdaLift)
+  {
+    // no lambdas are lazily lifted
+    return 0;
+  }
+  Trace("uf-ho") << "HoExtension::checkLazyLambda..." << std::endl;
+  NodeManager* nm = NodeManager::currentNM();
+  unsigned numLemmas = 0;
+  d_lambdaEqc.clear();
+  eq::EqualityEngine* ee = d_state.getEqualityEngine();
+  eq::EqClassesIterator eqcs_i = eq::EqClassesIterator(ee);
+  // normal functions equated to lambda functions
+  std::unordered_set<Node> normalEqFuns;
+  // mapping from functions to terms
+  while (!eqcs_i.isFinished())
+  {
+    Node eqc = (*eqcs_i);
+    ++eqcs_i;
+    if (!eqc.getType().isFunction())
+    {
+      continue;
+    }
+    eq::EqClassIterator eqc_i = eq::EqClassIterator(eqc, ee);
+    Node lamRep;  // the first lambda function we encounter in the equivalence
+                  // class
+    Node lamRepLam;
+    std::unordered_set<Node> normalEqFunWait;
+    while (!eqc_i.isFinished())
+    {
+      Node n = *eqc_i;
+      ++eqc_i;
+      Node lam = d_ll.getLambdaFor(n);
+      if (lam.isNull())
+      {
+        if (!lamRep.isNull())
+        {
+          // if we are equal to a lambda function, we must beta-reduce
+          // applications of this
+          normalEqFuns.insert(n);
+        }
+        else
+        {
+          // waiting to see if there is a lambda function in this equivalence
+          // class
+          normalEqFunWait.insert(n);
+        }
+      }
+      else if (lamRep.isNull())
+      {
+        // there is a lambda function in this equivalence class
+        lamRep = n;
+        lamRepLam = lam;
+        // must consider all normal functions we've seen so far
+        normalEqFuns.insert(normalEqFunWait.begin(), normalEqFunWait.end());
+        normalEqFunWait.clear();
+      }
+      else
+      {
+        // two lambda functions are in same equivalence class
+        Node f = lamRep < n ? lamRep : n;
+        Node g = lamRep < n ? n : lamRep;
+        Trace("uf-ho-debug") << "  found equivalent lambda functions " << f
+                             << " and " << g << std::endl;
+        Node flam = lamRep < n ? lamRepLam : lam;
+        Assert(!flam.isNull() && flam.getKind() == LAMBDA);
+        Node lhs = flam[1];
+        Node glam = lamRep < n ? lam : lamRepLam;
+        Trace("uf-ho-debug")
+            << "  lambda are " << flam << " and " << glam << std::endl;
+        std::vector<Node> args(flam[0].begin(), flam[0].end());
+        Node rhs = d_ll.betaReduce(glam, args);
+        Node univ = nm->mkNode(FORALL, flam[0], lhs.eqNode(rhs));
+        // f = g => forall x. reduce(lambda(f)(x)) = reduce(lambda(g)(x))
+        //
+        // For example, if f -> lambda z. z+1, g -> lambda y. y+3, this
+        // will infer: f = g => forall x. x+1 = x+3, which simplifies to
+        // f != g.
+        Node lem = nm->mkNode(IMPLIES, f.eqNode(g), univ);
+        if (cacheLemma(lem))
+        {
+          d_im.lemma(lem, InferenceId::UF_HO_LAMBDA_UNIV_EQ);
+          numLemmas++;
+        }
+      }
+    }
+    if (!lamRep.isNull())
+    {
+      d_lambdaEqc[eqc] = lamRep;
+    }
+  }
+  Trace("uf-ho-debug")
+      << "  found " << normalEqFuns.size()
+      << " ordinary functions that are equal to lambda functions" << std::endl;
+  if (normalEqFuns.empty())
+  {
+    return numLemmas;
+  }
+  // if we have normal functions that are equal to lambda functions, go back
+  // and ensure they are mapped properly
+  // mapping from functions to terms
+  eq::EqClassesIterator eqcs_i2 = eq::EqClassesIterator(ee);
+  while (!eqcs_i2.isFinished())
+  {
+    Node eqc = (*eqcs_i2);
+    ++eqcs_i2;
+    Trace("uf-ho-debug") << "Check equivalence class " << eqc << std::endl;
+    eq::EqClassIterator eqc_i = eq::EqClassIterator(eqc, ee);
+    while (!eqc_i.isFinished())
+    {
+      Node n = *eqc_i;
+      ++eqc_i;
+      Trace("uf-ho-debug") << "Check term " << n << std::endl;
+      Node op;
+      Kind k = n.getKind();
+      std::vector<Node> args;
+      if (k == APPLY_UF)
+      {
+        op = n.getOperator();
+        args.insert(args.end(), n.begin(), n.end());
+      }
+      else if (k == HO_APPLY)
+      {
+        op = n[0];
+        args.push_back(n[1]);
+      }
+      else
+      {
+        continue;
+      }
+      if (normalEqFuns.find(op) == normalEqFuns.end())
+      {
+        continue;
+      }
+      Trace("uf-ho-debug") << "  found relevant ordinary application " << n
+                           << std::endl;
+      Assert(ee->hasTerm(op));
+      Node r = ee->getRepresentative(op);
+      Assert(d_lambdaEqc.find(r) != d_lambdaEqc.end());
+      Node lf = d_lambdaEqc[r];
+      Node lam = d_ll.getLambdaFor(lf);
+      Assert(!lam.isNull() && lam.getKind() == LAMBDA);
+      // a normal function g equal to a lambda, say f --> lambda(f)
+      // need to infer f = g => g(t) = f(t) for all terms g(t)
+      // that occur in the equality engine.
+      Node premise = op.eqNode(lf);
+      args.insert(args.begin(), lam);
+      Node rhs = nm->mkNode(n.getKind(), args);
+      rhs = rewrite(rhs);
+      Node conc = n.eqNode(rhs);
+      Node lem = nm->mkNode(IMPLIES, premise, conc);
+      if (cacheLemma(lem))
+      {
+        d_im.lemma(lem, InferenceId::UF_HO_LAMBDA_APP_REDUCE);
+        numLemmas++;
+      }
+    }
+  }
+  return numLemmas;
+}
+
 unsigned HoExtension::check()
 {
   Trace("uf-ho") << "HoExtension::checkHigherOrder..." << std::endl;
@@ -429,14 +644,24 @@ unsigned HoExtension::check()
     }
   } while (num_facts > 0);
 
-  unsigned num_lemmas = 0;
-
-  num_lemmas = checkExtensionality();
-  if (num_lemmas > 0)
+  // Apply extensionality, lazy lambda schemas in order. We make lazy lambda
+  // handling come last as it may introduce quantifiers.
+  for (size_t i = 0; i < 2; i++)
   {
-    Trace("uf-ho") << "...extensionality returned " << num_lemmas << " lemmas."
-                   << std::endl;
-    return num_lemmas;
+    unsigned num_lemmas = 0;
+    // apply the schema
+    switch (i)
+    {
+      case 0: num_lemmas = checkExtensionality(); break;
+      case 1: num_lemmas = checkLazyLambda(); break;
+      default: break;
+    }
+    // finish if we added lemmas
+    if (num_lemmas > 0)
+    {
+      Trace("uf-ho") << "...returned " << num_lemmas << " lemmas." << std::endl;
+      return num_lemmas;
+    }
   }
 
   Trace("uf-ho") << "...finished check higher order." << std::endl;
@@ -464,6 +689,16 @@ bool HoExtension::collectModelInfoHo(TheoryModel* m,
   // non-standard alternative to using a type enumerator over function
   // values to assign unique values.
   int addedLemmas = checkExtensionality(m);
+  // for equivalence classes that we know to assign a lambda directly
+  for (const std::pair<const Node, Node>& p : d_lambdaEqc)
+  {
+    Node lam = d_ll.getLambdaFor(p.second);
+    Assert(!lam.isNull());
+    m->assertEquality(p.second, lam, true);
+    m->assertSkeleton(lam);
+    // assign it as the function definition for all variables in this class
+    m->assignFunctionDefinition(p.second, lam);
+  }
   return addedLemmas == 0;
 }
 
@@ -484,6 +719,17 @@ bool HoExtension::collectModelInfoHoTerm(Node n, TheoryModel* m)
   return true;
 }
 
+bool HoExtension::cacheLemma(TNode lem)
+{
+  Node rewritten = rewrite(lem);
+  if (d_cachedLemmas.find(rewritten) != d_cachedLemmas.end())
+  {
+    return false;
+  }
+  d_cachedLemmas.insert(rewritten);
+  return true;
+}
+
 }  // namespace uf
 }  // namespace theory
 }  // namespace cvc5
diff --git a/src/theory/uf/ho_extension.h b/src/theory/uf/ho_extension.h
index a646b57be..13d381622 100644
--- a/src/theory/uf/ho_extension.h
+++ b/src/theory/uf/ho_extension.h
@@ -15,14 +15,15 @@
 
 #include "cvc5_private.h"
 
-#ifndef __CVC5__THEORY__UF__HO_EXTENSION_H
-#define __CVC5__THEORY__UF__HO_EXTENSION_H
+#ifndef CVC5__THEORY__UF__HO_EXTENSION_H
+#define CVC5__THEORY__UF__HO_EXTENSION_H
 
 #include "context/cdhashmap.h"
 #include "context/cdhashset.h"
 #include "context/cdo.h"
 #include "expr/node.h"
 #include "smt/env_obj.h"
+#include "theory/skolem_lemma.h"
 #include "theory/theory_inference_manager.h"
 #include "theory/theory_model.h"
 #include "theory/theory_state.h"
@@ -31,7 +32,7 @@ namespace cvc5 {
 namespace theory {
 namespace uf {
 
-class TheoryUF;
+class LambdaLift;
 
 /** The higher-order extension of the theory of uninterpreted functions
  *
@@ -54,7 +55,10 @@ class HoExtension : protected EnvObj
   typedef context::CDHashMap<Node, Node> NodeNodeMap;
 
  public:
-  HoExtension(Env& env, TheoryState& state, TheoryInferenceManager& im);
+  HoExtension(Env& env,
+              TheoryState& state,
+              TheoryInferenceManager& im,
+              LambdaLift& ll);
 
   /** ppRewrite
    *
@@ -65,7 +69,7 @@ class HoExtension : protected EnvObj
    * function variables for function heads that are not variables via the
    * getApplyUfForHoApply method below.
    */
-  Node ppRewrite(Node node);
+  TrustNode ppRewrite(Node node, std::vector<SkolemLemma>& lems);
 
   /** check higher order
    *
@@ -171,8 +175,20 @@ class HoExtension : protected EnvObj
 
   /** check whether app-completion should be applied for any
    * pair of terms in the equality engine.
+   *
+   * Returns the number of lemmas added on this call.
    */
   unsigned checkAppCompletion();
+  /**
+   * Check lazy lambda.
+   *
+   * This assumes that lambdas are not eagerly lifted to quantified formulas.
+   * It processes two lemma schemas, UF_HO_LAMBDA_UNIV_EQ and
+   * UF_HO_LAMBDA_APP_REDUCE. For details on these, see inference_id.h.
+   *
+   * Returns the number of lemmas added on this call.
+   */
+  unsigned checkLazyLambda();
   /** collect model info for higher-order term
    *
    * This adds required constraints to m for term n. In particular, if n is
@@ -182,24 +198,43 @@ class HoExtension : protected EnvObj
   bool collectModelInfoHoTerm(Node n, TheoryModel* m);
 
  private:
+  /** Cache lemma lem, return true if it does not already exist */
+  bool cacheLemma(TNode lem);
   /** common constants */
   Node d_true;
   /** Reference to the state object */
   TheoryState& d_state;
   /** Reference to the inference manager */
   TheoryInferenceManager& d_im;
+  /** Lambda lifting utility */
+  LambdaLift& d_ll;
   /** extensionality has been applied to these disequalities */
   NodeSet d_extensionality;
+  /**
+   * The lemmas we have sent. This is required since the UF inference manager
+   * does not cache lemmas.
+   */
+  NodeSet d_cachedLemmas;
+  /**
+   * In the following, we say that a "lambda function" is a variable k that was
+   * introduced by the lambda lifting utility, and has a corresponding lambda
+   * definition.
+   *
+   * This maps equivalence class representatives that have lambda functions in
+   * them to one such lambda function. This map is computed at each full effort
+   * and valid only during collectModelInfoHo.
+   */
+  std::unordered_map<Node, Node> d_lambdaEqc;
 
   /** cache of getExtensionalityDeq below */
   std::map<Node, Node> d_extensionality_deq;
 
   /** map from non-standard operators to their skolems */
   NodeNodeMap d_uf_std_skolem;
-}; /* class TheoryUF */
+};
 
 }  // namespace uf
 }  // namespace theory
 }  // namespace cvc5
 
-#endif /* __CVC5__THEORY__UF__HO_EXTENSION_H */
+#endif
diff --git a/src/theory/uf/kinds b/src/theory/uf/kinds
index a1db5120f..252020b88 100644
--- a/src/theory/uf/kinds
+++ b/src/theory/uf/kinds
@@ -33,8 +33,6 @@ typerule LAMBDA ::cvc5::theory::uf::LambdaTypeRule
 
 variable BOOLEAN_TERM_VARIABLE "Boolean term variable"
 
-variable LAMBDA_VARIABLE "Lambda variable, used for lazy lambda lifting"
-
 # lambda expressions that are isomorphic to array constants can be considered constants
 construle LAMBDA ::cvc5::theory::uf::LambdaTypeRule
 
diff --git a/src/theory/uf/lambda_lift.cpp b/src/theory/uf/lambda_lift.cpp
index 3d51e6858..605e346af 100644
--- a/src/theory/uf/lambda_lift.cpp
+++ b/src/theory/uf/lambda_lift.cpp
@@ -90,6 +90,11 @@ Node LambdaLift::getLambdaFor(TNode skolem) const
   return it->second;
 }
 
+bool LambdaLift::isLambdaFunction(TNode n) const
+{
+  return !getLambdaFor(n).isNull();
+}
+
 Node LambdaLift::getAssertionFor(TNode node)
 {
   TNode skolem = getSkolemFor(node);
diff --git a/src/theory/uf/lambda_lift.h b/src/theory/uf/lambda_lift.h
index f61007d1c..1e36dad31 100644
--- a/src/theory/uf/lambda_lift.h
+++ b/src/theory/uf/lambda_lift.h
@@ -60,6 +60,8 @@ class LambdaLift : protected EnvObj
 
   /** Get the lambda term for skolem, if skolem is a lambda function. */
   Node getLambdaFor(TNode skolem) const;
+  /** Is skolem a lambda function? */
+  bool isLambdaFunction(TNode n) const;
 
   /**
    * Beta-reduce node. If node is APPLY_UF and its operator is a lambda
diff --git a/src/theory/uf/theory_uf.cpp b/src/theory/uf/theory_uf.cpp
index 5c9446507..b6d789773 100644
--- a/src/theory/uf/theory_uf.cpp
+++ b/src/theory/uf/theory_uf.cpp
@@ -21,6 +21,7 @@
 #include <sstream>
 
 #include "expr/node_algorithm.h"
+#include "expr/skolem_manager.h"
 #include "options/quantifiers_options.h"
 #include "options/smt_options.h"
 #include "options/theory_options.h"
@@ -31,6 +32,7 @@
 #include "theory/type_enumerator.h"
 #include "theory/uf/cardinality_extension.h"
 #include "theory/uf/ho_extension.h"
+#include "theory/uf/lambda_lift.h"
 #include "theory/uf/theory_uf_rewriter.h"
 
 using namespace std;
@@ -46,6 +48,7 @@ TheoryUF::TheoryUF(Env& env,
                    std::string instanceName)
     : Theory(THEORY_UF, env, out, valuation, instanceName),
       d_thss(nullptr),
+      d_lambdaLift(new LambdaLift(env)),
       d_ho(nullptr),
       d_functionsTerms(context()),
       d_symb(env, instanceName),
@@ -100,7 +103,7 @@ void TheoryUF::finishInit() {
   if (isHo)
   {
     d_equalityEngine->addFunctionKind(kind::HO_APPLY);
-    d_ho.reset(new HoExtension(d_env, d_state, d_im));
+    d_ho.reset(new HoExtension(d_env, d_state, d_im, *d_lambdaLift.get()));
   }
 }
 
@@ -212,38 +215,41 @@ TrustNode TheoryUF::ppRewrite(TNode node, std::vector<SkolemLemma>& lems)
   Trace("uf-exp-def") << "TheoryUF::ppRewrite: expanding definition : " << node
                       << std::endl;
   Kind k = node.getKind();
+  bool isHol = logicInfo().isHigherOrder();
   if (k == kind::HO_APPLY || (node.isVar() && node.getType().isFunction()))
   {
-    if (!logicInfo().isHigherOrder())
+    if (!isHol)
     {
       std::stringstream ss;
       ss << "Partial function applications are only supported with "
             "higher-order logic. Try adding the logic prefix HO_.";
       throw LogicException(ss.str());
     }
-    Node ret = d_ho->ppRewrite(node);
-    if (ret != node)
-    {
-      Trace("uf-exp-def") << "TheoryUF::ppRewrite: higher-order: " << node
-                          << " to " << ret << std::endl;
-      return TrustNode::mkTrustRewrite(node, ret, nullptr);
-    }
   }
   else if (k == kind::APPLY_UF)
   {
-    // check for higher-order
-    // logic exception if higher-order is not enabled
-    if (isHigherOrderType(node.getOperator().getType())
-        && !logicInfo().isHigherOrder())
+    if (!isHol && isHigherOrderType(node.getOperator().getType()))
     {
+      // check for higher-order
+      // logic exception if higher-order is not enabled
       std::stringstream ss;
       ss << "UF received an application whose operator has higher-order type "
          << node
-         << ", which is only supported with higher-order logic. Try adding the "
-            "logic prefix HO_.";
+         << ", which is only supported with higher-order logic. Try adding "
+            "the logic prefix HO_.";
       throw LogicException(ss.str());
     }
   }
+  if (isHol)
+  {
+    TrustNode ret = d_ho->ppRewrite(node, lems);
+    if (!ret.isNull())
+    {
+      Trace("uf-exp-def") << "TheoryUF::ppRewrite: higher-order: " << node
+                          << " to " << ret.getNode() << std::endl;
+      return ret;
+    }
+  }
   return TrustNode::null();
 }
 
@@ -308,6 +314,18 @@ void TheoryUF::preRegisterTerm(TNode node)
     d_equalityEngine->addTerm(node);
     break;
   }
+
+  if (logicInfo().isHigherOrder())
+  {
+    // When using lazy lambda handling, if node is a lambda function, it must
+    // be marked as a shared term. This is to ensure we split on the equality
+    // of lambda functions with other functions when doing care graph
+    // based theory combination.
+    if (d_lambdaLift->isLambdaFunction(node))
+    {
+      addSharedTerm(node);
+    }
+  }
 }
 
 void TheoryUF::explain(TNode literal, Node& exp)
@@ -523,7 +541,20 @@ bool TheoryUF::areCareDisequal(TNode x, TNode y)
     TNode y_shared =
         d_equalityEngine->getTriggerTermRepresentative(y, THEORY_UF);
     EqualityStatus eqStatus = d_valuation.getEqualityStatus(x_shared, y_shared);
-    if( eqStatus==EQUALITY_FALSE_AND_PROPAGATED || eqStatus==EQUALITY_FALSE || eqStatus==EQUALITY_FALSE_IN_MODEL ){
+    if (eqStatus == EQUALITY_FALSE || eqStatus == EQUALITY_FALSE_AND_PROPAGATED)
+    {
+      return true;
+    }
+    else if (eqStatus == EQUALITY_FALSE_IN_MODEL)
+    {
+      // if x or y is a lambda function, and they are neither entailed to
+      // be equal or disequal, then we return false. This ensures the pair
+      // (x,y) may be considered for the care graph.
+      if (d_lambdaLift->isLambdaFunction(x)
+          || d_lambdaLift->isLambdaFunction(y))
+      {
+        return false;
+      }
       return true;
     }
   }
diff --git a/src/theory/uf/theory_uf.h b/src/theory/uf/theory_uf.h
index 6f04035ed..b72262164 100644
--- a/src/theory/uf/theory_uf.h
+++ b/src/theory/uf/theory_uf.h
@@ -35,6 +35,7 @@ namespace uf {
 
 class CardinalityExtension;
 class HoExtension;
+class LambdaLift;
 
 class TheoryUF : public Theory {
  public:
@@ -74,6 +75,8 @@ class TheoryUF : public Theory {
 private:
   /** The associated cardinality extension (or nullptr if it does not exist) */
   std::unique_ptr<CardinalityExtension> d_thss;
+  /** the lambda lifting utility */
+  std::unique_ptr<LambdaLift> d_lambdaLift;
   /** the higher-order solver extension (or nullptr if it does not exist) */
   std::unique_ptr<HoExtension> d_ho;
 
diff --git a/test/regress/CMakeLists.txt b/test/regress/CMakeLists.txt
index d9076c81b..4937c8481 100644
--- a/test/regress/CMakeLists.txt
+++ b/test/regress/CMakeLists.txt
@@ -660,11 +660,14 @@ set(regress_0_tests
   regress0/ho/issue6536.smt2
   regress0/ho/ite-apply-eq.smt2
   regress0/ho/lambda-equality-non-canon.smt2
+  regress0/ho/lazy-lambda-model.smt2
   regress0/ho/match-middle.smt2
   regress0/ho/modulo-func-equality.smt2
   regress0/ho/qgu-fuzz-ho-1-dd.smt2
   regress0/ho/qgu-fuzz-ho-2-dd-no-ext.smt2
   regress0/ho/shadowing-defs.smt2
+  regress0/ho/simple-conf-lazy-lambda-lift.smt2
+  regress0/ho/simple-conf-lazy-lambda-lift-app.smt2
   regress0/ho/simple-matching-partial.smt2
   regress0/ho/simple-matching.smt2
   regress0/ho/trans.smt2
@@ -1604,6 +1607,7 @@ set(regress_1_tests
   regress1/bags/intersection_min1.smt2
   regress1/bags/intersection_min2.smt2
   regress1/bags/issue5759.smt2
+  regress1/bags/map-lazy-lam.smt2
   regress1/bags/map1.smt2
   regress1/bags/map2.smt2
   regress1/bags/map3.smt2
@@ -1728,8 +1732,10 @@ set(regress_1_tests
   regress1/fmf/with-ind-104-core.smt2
   regress1/gensys_brn001.smt2
   regress1/ho/bug_freeVar_BDD_General_data_270.p
+  regress1/ho/bug_freevar_PHI004^4-delta.smt2
   regress1/ho/bound_var_bug.p
   regress1/ho/fta0328.lfho.p
+  regress1/ho/ho-fun-sharing-dd.smt2
   regress1/ho/issue3136-fconst-bool-bool.smt2
   regress1/ho/issue4065-no-rep.smt2
   regress1/ho/issue4092-sinf.smt2
@@ -2805,8 +2811,6 @@ set(regression_disabled_tests
   ###
   regress1/bug472.smt2
   regress1/datatypes/non-simple-rec-set.smt2
-  # disabled temporarily until lazy lambda handling is merged
-  regress1/ho/bug_freevar_PHI004^4-delta.smt2
   # results in an assertion failure (see issue #1650).
   regress1/ho/hoa0102.smt2
   # after disallowing solving for terms with quantifiers
diff --git a/test/regress/regress0/ho/lazy-lambda-model.smt2 b/test/regress/regress0/ho/lazy-lambda-model.smt2
new file mode 100644
index 000000000..71df017e9
--- /dev/null
+++ b/test/regress/regress0/ho/lazy-lambda-model.smt2
@@ -0,0 +1,16 @@
+; COMMAND-LINE: --uf-lazy-ll
+; EXPECT: sat
+(set-logic HO_ALL)
+(declare-fun g (Int) Int)
+(declare-fun h (Int) Int)
+(define-fun f ((x Int)) Int (ite (> x 0) (* 2 x) x))
+
+(declare-fun a () Int)
+(declare-fun b () Int)
+
+
+(assert (or (= f g) (= f h)))
+
+(assert (and (= (h a) 26) (= (g b) 26)))
+
+(check-sat)
diff --git a/test/regress/regress0/ho/simple-conf-lazy-lambda-lift-app.smt2 b/test/regress/regress0/ho/simple-conf-lazy-lambda-lift-app.smt2
new file mode 100644
index 000000000..ab949c40c
--- /dev/null
+++ b/test/regress/regress0/ho/simple-conf-lazy-lambda-lift-app.smt2
@@ -0,0 +1,14 @@
+; COMMAND-LINE: --uf-lazy-ll
+; EXPECT: unsat
+(set-logic HO_ALL)
+(set-info :status unsat)
+(define-fun f ((x Int)) Int (+ x 1))
+(declare-fun g (Int) Int)
+(declare-fun h (Int) Int)
+
+(assert (or (= f g) (= f h)))
+
+(assert (= (g 4) 0))
+(assert (= (h 4) 8))
+
+(check-sat)
diff --git a/test/regress/regress0/ho/simple-conf-lazy-lambda-lift.smt2 b/test/regress/regress0/ho/simple-conf-lazy-lambda-lift.smt2
new file mode 100644
index 000000000..d11a4418b
--- /dev/null
+++ b/test/regress/regress0/ho/simple-conf-lazy-lambda-lift.smt2
@@ -0,0 +1,11 @@
+; COMMAND-LINE: --uf-lazy-ll
+; EXPECT: unsat
+(set-logic HO_ALL)
+(set-info :status unsat)
+(define-fun f ((x Int)) Int (+ x 1))
+(define-fun g ((x Int)) Int (+ x 2))
+(define-fun h ((x Int)) Int 6)
+
+(assert (or (= f g) (= f h)))
+
+(check-sat)
diff --git a/test/regress/regress1/bags/map-lazy-lam.smt2 b/test/regress/regress1/bags/map-lazy-lam.smt2
new file mode 100644
index 000000000..c99bab7c9
--- /dev/null
+++ b/test/regress/regress1/bags/map-lazy-lam.smt2
@@ -0,0 +1,9 @@
+; COMMAND-LINE: --fmf-bound --uf-lazy-ll
+; EXPECT: sat
+(set-logic HO_ALL)
+(define-fun f ((x String)) Int 1)
+(define-fun cardinality ((A (Bag String))) Int
+  (bag.count 1 (bag.map f A)))
+(declare-fun A () (Bag String))
+(assert (= (cardinality A) 20))
+(check-sat)
diff --git a/test/regress/regress1/ho/bug_freevar_PHI004^4-delta.smt2 b/test/regress/regress1/ho/bug_freevar_PHI004^4-delta.smt2
index 7c9de4c49..0155eecb2 100644
--- a/test/regress/regress1/ho/bug_freevar_PHI004^4-delta.smt2
+++ b/test/regress/regress1/ho/bug_freevar_PHI004^4-delta.smt2
@@ -1,4 +1,4 @@
-; COMMAND-LINE: --finite-model-find -q --decision=justification-old
+; COMMAND-LINE: --finite-model-find --decision=justification-old --uf-lazy-ll -q
 ; EXPECT: sat
 
 (set-logic HO_ALL)
diff --git a/test/regress/regress1/ho/ho-fun-sharing-dd.smt2 b/test/regress/regress1/ho/ho-fun-sharing-dd.smt2
new file mode 100644
index 000000000..300dc26de
--- /dev/null
+++ b/test/regress/regress1/ho/ho-fun-sharing-dd.smt2
@@ -0,0 +1,10 @@
+; COMMAND-LINE: --finite-model-find
+; EXPECT: unsat
+(set-logic HO_ALL)
+(declare-sort n 0)
+(declare-fun x () n)
+(declare-fun s (n) n)
+(declare-fun t ((-> n Bool)) Bool)
+(assert (forall ((X n)) (t (lambda ((Xu n)) (= X (s Xu))))))
+(assert (not (t (lambda ((Xu n)) (= x (s Xu))))))
+(check-sat)