1 files changed, 47 insertions, 16 deletions

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;
     }
   }