Split higher-order UF solver (#2890)

5 files changed, 658 insertions, 423 deletions

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index f79c82d49..3b0e14976 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -667,6 +667,8 @@ libcvc4_add_sources(
   theory/uf/equality_engine.cpp
   theory/uf/equality_engine.h
   theory/uf/equality_engine_types.h
+  theory/uf/ho_extension.cpp
+  theory/uf/ho_extension.h
   theory/uf/symmetry_breaker.cpp
   theory/uf/symmetry_breaker.h
   theory/uf/theory_uf.cpp
diff --git a/src/theory/uf/ho_extension.cpp b/src/theory/uf/ho_extension.cpp
new file mode 100644
index 000000000..4d5a879ad
--- /dev/null
+++ b/src/theory/uf/ho_extension.cpp
@@ -0,0 +1,433 @@
+/*********************                                                        */
+/*! \file ho_extension.cpp
+ ** \verbatim
+ ** Top contributors (to current version):
+ **   Andrew Reynolds
+ ** This file is part of the CVC4 project.
+ ** Copyright (c) 2009-2019 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.\endverbatim
+ **
+ ** \brief Implementation of the higher-order extension of TheoryUF.
+ **
+ **/
+
+#include "theory/uf/ho_extension.h"
+
+#include "expr/node_algorithm.h"
+#include "options/uf_options.h"
+#include "theory/theory_model.h"
+#include "theory/uf/theory_uf.h"
+#include "theory/uf/theory_uf_rewriter.h"
+
+using namespace std;
+using namespace CVC4::kind;
+
+namespace CVC4 {
+namespace theory {
+namespace uf {
+
+HoExtension::HoExtension(TheoryUF& p,
+                         context::Context* c,
+                         context::UserContext* u)
+    : d_parent(p), d_extensionality(u), d_uf_std_skolem(u)
+{
+  d_true = NodeManager::currentNM()->mkConst(true);
+}
+
+Node HoExtension::expandDefinition(Node node)
+{
+  // 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 : expandDefinition : " << node << " to " << ret
+                   << std::endl;
+    return ret;
+  }
+  return node;
+}
+
+Node HoExtension::getExtensionalityDeq(TNode deq)
+{
+  Assert(deq.getKind() == NOT && deq[0].getKind() == EQUAL);
+  Assert(deq[0][0].getType().isFunction());
+  std::map<Node, Node>::iterator it = d_extensionality_deq.find(deq);
+  if (it == d_extensionality_deq.end())
+  {
+    TypeNode tn = deq[0][0].getType();
+    std::vector<TypeNode> argTypes = tn.getArgTypes();
+    std::vector<Node> skolems;
+    NodeManager* nm = NodeManager::currentNM();
+    for (unsigned i = 0, nargs = argTypes.size(); i < nargs; i++)
+    {
+      Node k =
+          nm->mkSkolem("k", argTypes[i], "skolem created for extensionality.");
+      skolems.push_back(k);
+    }
+    Node t[2];
+    for (unsigned i = 0; i < 2; i++)
+    {
+      std::vector<Node> children;
+      Node curr = deq[0][i];
+      while (curr.getKind() == HO_APPLY)
+      {
+        children.push_back(curr[1]);
+        curr = curr[0];
+      }
+      children.push_back(curr);
+      std::reverse(children.begin(), children.end());
+      children.insert(children.end(), skolems.begin(), skolems.end());
+      t[i] = nm->mkNode(APPLY_UF, children);
+    }
+    Node conc = t[0].eqNode(t[1]).negate();
+    d_extensionality_deq[deq] = conc;
+    return conc;
+  }
+  return it->second;
+}
+
+unsigned HoExtension::applyExtensionality(TNode deq)
+{
+  Assert(deq.getKind() == NOT && deq[0].getKind() == EQUAL);
+  Assert(deq[0][0].getType().isFunction());
+  // apply extensionality
+  if (d_extensionality.find(deq) == d_extensionality.end())
+  {
+    d_extensionality.insert(deq);
+    Node conc = getExtensionalityDeq(deq);
+    Node lem = NodeManager::currentNM()->mkNode(OR, deq[0], conc);
+    Trace("uf-ho-lemma") << "uf-ho-lemma : extensionality : " << lem
+                         << std::endl;
+    d_parent.getOutputChannel().lemma(lem);
+    return 1;
+  }
+  return 0;
+}
+
+Node HoExtension::getApplyUfForHoApply(Node node)
+{
+  Assert(node[0].getType().getNumChildren() == 2);
+  std::vector<TNode> args;
+  Node f = TheoryUfRewriter::decomposeHoApply(node, args, true);
+  Node new_f = f;
+  NodeManager* nm = NodeManager::currentNM();
+  if (!TheoryUfRewriter::canUseAsApplyUfOperator(f))
+  {
+    NodeNodeMap::const_iterator itus = d_uf_std_skolem.find(f);
+    if (itus == d_uf_std_skolem.end())
+    {
+      std::unordered_set<Node, NodeHashFunction> fvs;
+      expr::getFreeVariables(f, fvs);
+      Node lem;
+      if (!fvs.empty())
+      {
+        std::vector<TypeNode> newTypes;
+        std::vector<Node> vs;
+        std::vector<Node> nvs;
+        for (const Node& v : fvs)
+        {
+          TypeNode vt = v.getType();
+          newTypes.push_back(vt);
+          Node nv = nm->mkBoundVar(vt);
+          vs.push_back(v);
+          nvs.push_back(nv);
+        }
+        TypeNode ft = f.getType();
+        std::vector<TypeNode> argTypes = ft.getArgTypes();
+        TypeNode rangeType = ft.getRangeType();
+
+        newTypes.insert(newTypes.end(), argTypes.begin(), argTypes.end());
+        TypeNode nft = nm->mkFunctionType(newTypes, rangeType);
+        new_f = nm->mkSkolem("app_uf", nft);
+        for (const Node& v : vs)
+        {
+          new_f = nm->mkNode(HO_APPLY, new_f, v);
+        }
+        Assert(new_f.getType() == f.getType());
+        Node eq = new_f.eqNode(f);
+        Node seq = eq.substitute(vs.begin(), vs.end(), nvs.begin(), nvs.end());
+        lem = nm->mkNode(
+            FORALL, nm->mkNode(BOUND_VAR_LIST, nvs), seq);
+      }
+      else
+      {
+        // introduce skolem to make a standard APPLY_UF
+        new_f = nm->mkSkolem("app_uf", f.getType());
+        lem = new_f.eqNode(f);
+      }
+      Trace("uf-ho-lemma")
+          << "uf-ho-lemma : Skolem definition for apply-conversion : " << lem
+          << std::endl;
+      d_parent.getOutputChannel().lemma(lem);
+      d_uf_std_skolem[f] = new_f;
+    }
+    else
+    {
+      new_f = (*itus).second;
+    }
+    // unroll the HO_APPLY, adding to the first argument position
+    // Note arguments in the vector args begin at position 1.
+    while (new_f.getKind() == HO_APPLY)
+    {
+      args.insert(args.begin() + 1, new_f[1]);
+      new_f = new_f[0];
+    }
+  }
+  Assert(TheoryUfRewriter::canUseAsApplyUfOperator(new_f));
+  args[0] = new_f;
+  Node ret = nm->mkNode(APPLY_UF, args);
+  Assert(ret.getType() == node.getType());
+  return ret;
+}
+
+unsigned HoExtension::checkExtensionality(TheoryModel* m)
+{
+  eq::EqualityEngine* ee = d_parent.getEqualityEngine();
+  unsigned num_lemmas = 0;
+  bool isCollectModel = (m != nullptr);
+  Trace("uf-ho") << "HoExtension::checkExtensionality, collectModel="
+                 << isCollectModel << "..." << std::endl;
+  std::map<TypeNode, std::vector<Node> > func_eqcs;
+  eq::EqClassesIterator eqcs_i = eq::EqClassesIterator(ee);
+  while (!eqcs_i.isFinished())
+  {
+    Node eqc = (*eqcs_i);
+    TypeNode tn = eqc.getType();
+    if (tn.isFunction())
+    {
+      // if during collect model, must have an infinite type
+      // if not during collect model, must have a finite type
+      if (tn.isInterpretedFinite() != isCollectModel)
+      {
+        func_eqcs[tn].push_back(eqc);
+        Trace("uf-ho-debug")
+            << "  func eqc : " << tn << " : " << eqc << std::endl;
+      }
+    }
+    ++eqcs_i;
+  }
+
+  for (std::map<TypeNode, std::vector<Node> >::iterator itf = func_eqcs.begin();
+       itf != func_eqcs.end();
+       ++itf)
+  {
+    for (unsigned j = 0, sizej = itf->second.size(); j < sizej; j++)
+    {
+      for (unsigned k = (j + 1), sizek = itf->second.size(); k < sizek; k++)
+      {
+        // if these equivalence classes are not explicitly disequal, do
+        // extensionality to ensure distinctness
+        if (!ee->areDisequal(itf->second[j], itf->second[k], false))
+        {
+          Node deq =
+              Rewriter::rewrite(itf->second[j].eqNode(itf->second[k]).negate());
+          // either add to model, or add lemma
+          if (isCollectModel)
+          {
+            // add extentionality disequality to the model
+            Node edeq = getExtensionalityDeq(deq);
+            Assert(edeq.getKind() == NOT && edeq[0].getKind() == EQUAL);
+            // introducing terms, must add required constraints, e.g. to
+            // force equalities between APPLY_UF and HO_APPLY terms
+            for (unsigned r = 0; r < 2; r++)
+            {
+              if (!collectModelInfoHoTerm(edeq[0][r], m))
+              {
+                return 1;
+              }
+            }
+            Trace("uf-ho-debug")
+                << "Add extensionality deq to model : " << edeq << std::endl;
+            if (!m->assertEquality(edeq[0][0], edeq[0][1], false))
+            {
+              return 1;
+            }
+          }
+          else
+          {
+            // apply extensionality lemma
+            num_lemmas += applyExtensionality(deq);
+          }
+        }
+      }
+    }
+  }
+  return num_lemmas;
+}
+
+unsigned HoExtension::applyAppCompletion(TNode n)
+{
+  Assert(n.getKind() == APPLY_UF);
+
+  eq::EqualityEngine* ee = d_parent.getEqualityEngine();
+  // must expand into APPLY_HO version if not there already
+  Node ret = TheoryUfRewriter::getHoApplyForApplyUf(n);
+  if (!ee->hasTerm(ret) || !ee->areEqual(ret, n))
+  {
+    Node eq = ret.eqNode(n);
+    Trace("uf-ho-lemma") << "uf-ho-lemma : infer, by apply-expand : " << eq
+                         << std::endl;
+    ee->assertEquality(eq, true, d_true);
+    return 1;
+  }
+  Trace("uf-ho-debug") << "    ...already have " << ret << " == " << n << "."
+                       << std::endl;
+  return 0;
+}
+
+unsigned HoExtension::checkAppCompletion()
+{
+  Trace("uf-ho") << "HoExtension::checkApplyCompletion..." << std::endl;
+  // compute the operators that are relevant (those for which an HO_APPLY exist)
+  std::set<TNode> rlvOp;
+  eq::EqualityEngine* ee = d_parent.getEqualityEngine();
+  eq::EqClassesIterator eqcs_i = eq::EqClassesIterator(ee);
+  std::map<TNode, std::vector<Node> > apply_uf;
+  while (!eqcs_i.isFinished())
+  {
+    Node eqc = (*eqcs_i);
+    Trace("uf-ho-debug") << "  apply completion : visit eqc " << eqc
+                         << std::endl;
+    eq::EqClassIterator eqc_i = eq::EqClassIterator(eqc, ee);
+    while (!eqc_i.isFinished())
+    {
+      Node n = *eqc_i;
+      if (n.getKind() == APPLY_UF || n.getKind() == HO_APPLY)
+      {
+        int curr_sum = 0;
+        std::map<TNode, bool> curr_rops;
+        if (n.getKind() == APPLY_UF)
+        {
+          TNode rop = ee->getRepresentative(n.getOperator());
+          if (rlvOp.find(rop) != rlvOp.end())
+          {
+            // try if its operator is relevant
+            curr_sum = applyAppCompletion(n);
+            if (curr_sum > 0)
+            {
+              return curr_sum;
+            }
+          }
+          else
+          {
+            // add to pending list
+            apply_uf[rop].push_back(n);
+          }
+          // arguments are also relevant operators  FIXME (github issue #1115)
+          for (unsigned k = 0; k < n.getNumChildren(); k++)
+          {
+            if (n[k].getType().isFunction())
+            {
+              TNode rop = ee->getRepresentative(n[k]);
+              curr_rops[rop] = true;
+            }
+          }
+        }
+        else
+        {
+          Assert(n.getKind() == HO_APPLY);
+          TNode rop = ee->getRepresentative(n[0]);
+          curr_rops[rop] = true;
+        }
+        for (std::map<TNode, bool>::iterator itc = curr_rops.begin();
+             itc != curr_rops.end();
+             ++itc)
+        {
+          TNode rop = itc->first;
+          if (rlvOp.find(rop) == rlvOp.end())
+          {
+            rlvOp.insert(rop);
+            // now, try each pending APPLY_UF for this operator
+            std::map<TNode, std::vector<Node> >::iterator itu =
+                apply_uf.find(rop);
+            if (itu != apply_uf.end())
+            {
+              for (unsigned j = 0, size = itu->second.size(); j < size; j++)
+              {
+                curr_sum = applyAppCompletion(itu->second[j]);
+                if (curr_sum > 0)
+                {
+                  return curr_sum;
+                }
+              }
+            }
+          }
+        }
+      }
+      ++eqc_i;
+    }
+    ++eqcs_i;
+  }
+  return 0;
+}
+
+unsigned HoExtension::check()
+{
+  Trace("uf-ho") << "HoExtension::checkHigherOrder..." << std::endl;
+
+  // infer new facts based on apply completion until fixed point
+  unsigned num_facts;
+  do
+  {
+    num_facts = checkAppCompletion();
+    if (d_parent.inConflict())
+    {
+      Trace("uf-ho") << "...conflict during app-completion." << std::endl;
+      return 1;
+    }
+  } while (num_facts > 0);
+
+  if (options::ufHoExt())
+  {
+    unsigned num_lemmas = 0;
+
+    num_lemmas = checkExtensionality();
+    if (num_lemmas > 0)
+    {
+      Trace("uf-ho") << "...extensionality returned " << num_lemmas
+                     << " lemmas." << std::endl;
+      return num_lemmas;
+    }
+  }
+
+  Trace("uf-ho") << "...finished check higher order." << std::endl;
+
+  return 0;
+}
+
+bool HoExtension::collectModelInfoHo(std::set<Node>& termSet, TheoryModel* m)
+{
+  for (std::set<Node>::iterator it = termSet.begin(); it != termSet.end(); ++it)
+  {
+    Node n = *it;
+    // For model-building with ufHo, we require that APPLY_UF is always
+    // expanded to HO_APPLY. That is, we always expand to a fully applicative
+    // encoding during model construction.
+    if (!collectModelInfoHoTerm(n, m))
+    {
+      return false;
+    }
+  }
+  int addedLemmas = checkExtensionality(m);
+  return addedLemmas == 0;
+}
+
+bool HoExtension::collectModelInfoHoTerm(Node n, TheoryModel* m)
+{
+  if (n.getKind() == APPLY_UF)
+  {
+    Node hn = TheoryUfRewriter::getHoApplyForApplyUf(n);
+    if (!m->assertEquality(n, hn, true))
+    {
+      return false;
+    }
+  }
+  return true;
+}
+
+}  // namespace uf
+}  // namespace theory
+}  // namespace CVC4
diff --git a/src/theory/uf/ho_extension.h b/src/theory/uf/ho_extension.h
new file mode 100644
index 000000000..a38d1803a
--- /dev/null
+++ b/src/theory/uf/ho_extension.h
@@ -0,0 +1,195 @@
+/*********************                                                        */
+/*! \file ho_extension.h
+ ** \verbatim
+ ** Top contributors (to current version):
+ **   Andrew Reynolds
+ ** This file is part of the CVC4 project.
+ ** Copyright (c) 2009-2019 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.\endverbatim
+ **
+ ** \brief The higher-order extension of TheoryUF.
+ **/
+
+#include "cvc4_private.h"
+
+#ifndef __CVC4__THEORY__UF__HO_EXTENSION_H
+#define __CVC4__THEORY__UF__HO_EXTENSION_H
+
+#include "context/cdhashmap.h"
+#include "context/cdhashset.h"
+#include "context/cdo.h"
+#include "expr/node.h"
+#include "theory/theory_model.h"
+
+namespace CVC4 {
+namespace theory {
+namespace uf {
+
+class TheoryUF;
+
+/** The higher-order extension of the theory of uninterpreted functions
+ *
+ * This extension is capable of handling UF constraints involving partial
+ * applications via the applicative encoding (kind HO_APPLY), and
+ * (dis)equalities involving function sorts. It uses a lazy applicative
+ * encoding and implements two axiom schemes, at a high level:
+ *
+ * (1) Extensionality, where disequalities between functions witnessed by
+ * arguments where the two functions differ,
+ *
+ * (2) App-Encode, where full applications of UF (kind APPLY_UF) are equated
+ * with curried applications (kind HO_APPLY).
+ *
+ * For more details, see "Extending SMT Solvers to Higher-Order", Barbosa et al.
+ */
+class HoExtension
+{
+  typedef context::CDHashSet<Node, NodeHashFunction> NodeSet;
+  typedef context::CDHashMap<Node, Node, NodeHashFunction> NodeNodeMap;
+
+ public:
+  HoExtension(TheoryUF& p, context::Context* c, context::UserContext* u);
+
+  /** expand definition
+   *
+   * This returns the expanded form of node.
+   *
+   * In particular, this function will convert applications of HO_APPLY
+   * to applications of APPLY_UF if they are fully applied, and introduce
+   * function variables for function heads that are not variables via the
+   * getApplyUfForHoApply method below.
+   */
+  Node expandDefinition(Node node);
+
+  /** check higher order
+   *
+   * This is called at full effort and infers facts and sends lemmas
+   * based on higher-order reasoning (specifically, extentionality and
+   * app completion). It returns the number of lemmas plus facts
+   * added to the equality engine.
+   */
+  unsigned check();
+
+  /** applyExtensionality
+   *
+   * Given disequality deq f != g, if not already cached, this sends a lemma of
+   * the form:
+   *   f = g V (f k) != (g k) for fresh constant k on the output channel of the
+   * parent TheoryUF object. This is an "extensionality lemma".
+   * Return value is the number of lemmas of this form sent on the output
+   * channel.
+   */
+  unsigned applyExtensionality(TNode deq);
+
+  /** collect model info
+   *
+   * This method adds the necessary equalities to the model m such that
+   * model construction is possible if this method returns true. These
+   * equalities may include HO_APPLY versions of all APPLY_UF terms.
+   *
+   * The argument termSet is the set of relevant terms that the parent TheoryUF
+   * object has added to m that belong to TheoryUF.
+   *
+   * This method ensures that the function variables in termSet
+   * respect extensionality. If some pair does not, then this method adds an
+   * extensionality equality directly to the equality engine of m.
+   *
+   * In more detail, functions f and g do not respect extensionality if f and g
+   * are not equal in the model, and there is not a pair of unequal witness
+   * terms f(k), g(k). In this case, we add the disequality
+   *    f(k') != g(k')
+   * for fresh (tuple) of variables k' to the equality engine of m. Notice
+   * this is done only for functions whose type has infinite cardinality,
+   * since all functions with finite cardinality are ensured to respect
+   * extensionality by this point due to our extentionality inference schema.
+   *
+   * If this method returns true, then all pairs of functions that are in
+   * distinct equivalence classes will be guaranteed to be assigned different
+   * values in m. It returns false if any (dis)equality added to m led to
+   * an inconsistency in m.
+   */
+  bool collectModelInfoHo(std::set<Node>& termSet, TheoryModel* m);
+
+ protected:
+  /** get apply uf for ho apply
+   *
+   * This returns the APPLY_UF equivalent for the HO_APPLY term node, where
+   * node has non-functional return type (that is, it corresponds to a fully
+   * applied function term).
+   * This call may introduce a skolem for the head operator and send out a lemma
+   * specifying the definition.
+   */
+  Node getApplyUfForHoApply(Node node);
+
+  /** get extensionality disequality
+   *
+   * Given disequality deq f != g, this returns the disequality:
+   *   (f k) != (g k) for fresh constant(s) k.
+   */
+  Node getExtensionalityDeq(TNode deq);
+
+  /**
+   * Check whether extensionality should be applied for any pair of terms in the
+   * equality engine.
+   *
+   * If we pass a null model m to this function, then we add extensionality
+   * lemmas to the output channel and return the total number of lemmas added.
+   * We only add lemmas for functions whose type is finite, since pairs of
+   * functions whose types are infinite can be made disequal in a model by
+   * witnessing a point they are disequal.
+   *
+   * If we pass a non-null model m to this function, then we add disequalities
+   * that correspond to the conclusion of extensionality lemmas to the model's
+   * equality engine. We return 0 if the equality engine of m is consistent
+   * after this call, and 1 otherwise. We only add disequalities for functions
+   * whose type is infinite, since our decision procedure guarantees that
+   * extensionality lemmas are added for all pairs of functions whose types are
+   * finite.
+   */
+  unsigned checkExtensionality(TheoryModel* m = nullptr);
+
+  /** applyAppCompletion
+   * This infers a correspondence between APPLY_UF and HO_APPLY
+   * versions of terms for higher-order.
+   * Given APPLY_UF node e.g. (f a b c), this adds the equality to its
+   * HO_APPLY equivalent:
+   *   (f a b c) == (@ (@ (@ f a) b) c)
+   * to equality engine, if not already equal.
+   * Return value is the number of equalities added.
+   */
+  unsigned applyAppCompletion(TNode n);
+
+  /** check whether app-completion should be applied for any
+   * pair of terms in the equality engine.
+   */
+  unsigned checkAppCompletion();
+  /** collect model info for higher-order term
+   *
+   * This adds required constraints to m for term n. In particular, if n is
+   * an APPLY_UF term, we add its HO_APPLY equivalent in this call. We return
+   * true if the model m is consistent after this call.
+   */
+  bool collectModelInfoHoTerm(Node n, TheoryModel* m);
+
+ private:
+  /** common constants */
+  Node d_true;
+  /** the parent of this extension */
+  TheoryUF& d_parent;
+  /** extensionality has been applied to these disequalities */
+  NodeSet d_extensionality;
+
+  /** 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 CVC4
+
+#endif /* __CVC4__THEORY__UF__HO_EXTENSION_H */
diff --git a/src/theory/uf/theory_uf.cpp b/src/theory/uf/theory_uf.cpp
index 508bd5002..393d9f640 100644
--- a/src/theory/uf/theory_uf.cpp
+++ b/src/theory/uf/theory_uf.cpp
@@ -29,6 +29,7 @@
 #include "proof/uf_proof.h"
 #include "theory/theory_model.h"
 #include "theory/type_enumerator.h"
+#include "theory/uf/ho_extension.h"
 #include "theory/uf/theory_uf_rewriter.h"
 #include "theory/uf/theory_uf_strong_solver.h"
 
@@ -52,8 +53,6 @@ TheoryUF::TheoryUF(context::Context* c,
       d_thss(NULL),
       d_equalityEngine(d_notify, c, instanceName + "theory::uf::ee", true),
       d_conflict(c, false),
-      d_extensionality(u),
-      d_uf_std_skolem(u),
       d_functionsTerms(c),
       d_symb(u, instanceName)
 {
@@ -61,9 +60,6 @@ TheoryUF::TheoryUF(context::Context* c,
 
   // The kinds we are treating as function application in congruence
   d_equalityEngine.addFunctionKind(kind::APPLY_UF, false, options::ufHo());
-  if( options::ufHo() ){
-    d_equalityEngine.addFunctionKind(kind::HO_APPLY);
-  }
 }
 
 TheoryUF::~TheoryUF() {
@@ -87,6 +83,11 @@ void TheoryUF::finishInit() {
   {
     d_thss = new StrongSolverTheoryUF(getSatContext(), getUserContext(), *d_out, this);
   }
+  if (options::ufHo())
+  {
+    d_equalityEngine.addFunctionKind(kind::HO_APPLY);
+    d_ho = new HoExtension(*this, getSatContext(), getUserContext());
+  }
 }
 
 static Node mkAnd(const std::vector<TNode>& conjunctions) {
@@ -143,7 +144,8 @@ void TheoryUF::check(Effort level) {
       d_equalityEngine.assertEquality(atom, polarity, fact);
       if( options::ufHo() && options::ufHoExt() ){
         if( !polarity && !d_conflict && atom[0].getType().isFunction() ){
-          applyExtensionality( fact );
+          // apply extensionality eagerly using the ho extension
+          d_ho->applyExtensionality(fact);
         }
       }
     } else if (atom.getKind() == kind::CARDINALITY_CONSTRAINT || atom.getKind() == kind::COMBINED_CARDINALITY_CONSTRAINT) {
@@ -168,89 +170,22 @@ void TheoryUF::check(Effort level) {
   }
 
   if(! d_conflict ){
+    // check with the cardinality constraints extension
     if (d_thss != NULL) {
       d_thss->check(level);
       if( d_thss->isConflict() ){
         d_conflict = true;
       }
     }
+    // check with the higher-order extension
     if(! d_conflict && fullEffort(level) ){
       if( options::ufHo() ){
-        checkHigherOrder();
+        d_ho->check();
       }
     }
   }
 }/* TheoryUF::check() */
 
-Node TheoryUF::getApplyUfForHoApply( Node node ) {
-  Assert( node[0].getType().getNumChildren()==2 );
-  std::vector< TNode > args;
-  Node f = TheoryUfRewriter::decomposeHoApply( node, args, true );
-  Node new_f = f;
-  NodeManager* nm = NodeManager::currentNM();
-  if( !TheoryUfRewriter::canUseAsApplyUfOperator( f ) ){
-    NodeNodeMap::const_iterator itus = d_uf_std_skolem.find( f );
-    if( itus==d_uf_std_skolem.end() ){
-      std::unordered_set<Node, NodeHashFunction> fvs;
-      expr::getFreeVariables(f, fvs);
-      Node lem;
-      if (!fvs.empty())
-      {
-        std::vector<TypeNode> newTypes;
-        std::vector<Node> vs;
-        std::vector<Node> nvs;
-        for (const Node& v : fvs)
-        {
-          TypeNode vt = v.getType();
-          newTypes.push_back(vt);
-          Node nv = nm->mkBoundVar(vt);
-          vs.push_back(v);
-          nvs.push_back(nv);
-        }
-        TypeNode ft = f.getType();
-        std::vector<TypeNode> argTypes = ft.getArgTypes();
-        TypeNode rangeType = ft.getRangeType();
-
-        newTypes.insert(newTypes.end(), argTypes.begin(), argTypes.end());
-        TypeNode nft = nm->mkFunctionType(newTypes, rangeType);
-        new_f = nm->mkSkolem("app_uf", nft);
-        for (const Node& v : vs)
-        {
-          new_f = nm->mkNode(kind::HO_APPLY, new_f, v);
-        }
-        Assert(new_f.getType() == f.getType());
-        Node eq = new_f.eqNode(f);
-        Node seq = eq.substitute(vs.begin(), vs.end(), nvs.begin(), nvs.end());
-        lem = nm->mkNode(
-            kind::FORALL, nm->mkNode(kind::BOUND_VAR_LIST, nvs), seq);
-      }
-      else
-      {
-        // introduce skolem to make a standard APPLY_UF
-        new_f = nm->mkSkolem("app_uf", f.getType());
-        lem = new_f.eqNode(f);
-      }
-      Trace("uf-ho-lemma") << "uf-ho-lemma : Skolem definition for apply-conversion : " << lem << std::endl;
-      d_out->lemma( lem );
-      d_uf_std_skolem[f] = new_f;
-    }else{
-      new_f = (*itus).second;
-    }
-    // unroll the HO_APPLY, adding to the first argument position
-    // Note arguments in the vector args begin at position 1.
-    while (new_f.getKind() == kind::HO_APPLY)
-    {
-      args.insert(args.begin() + 1, new_f[1]);
-      new_f = new_f[0];
-    }
-  }
-  Assert( TheoryUfRewriter::canUseAsApplyUfOperator( new_f ) );
-  args[0] = new_f;
-  Node ret = nm->mkNode(kind::APPLY_UF, args);
-  Assert(ret.getType() == node.getType());
-  return ret;
-}
-
 Node TheoryUF::getOperatorForApplyTerm( TNode node ) {
   Assert( node.getKind()==kind::APPLY_UF || node.getKind()==kind::HO_APPLY );
   if( node.getKind()==kind::APPLY_UF ){
@@ -266,18 +201,19 @@ unsigned TheoryUF::getArgumentStartIndexForApplyTerm( TNode node ) {
 }
 
 Node TheoryUF::expandDefinition(LogicRequest &logicRequest, Node node) {
-  Trace("uf-ho-debug") << "uf-ho-debug : expanding definition : " << node << std::endl;
+  Trace("uf-exp-def") << "TheoryUF::expandDefinition: expanding definition : "
+                      << node << std::endl;
   if( node.getKind()==kind::HO_APPLY ){
     if( !options::ufHo() ){
       std::stringstream ss;
       ss << "Partial function applications are not supported in default mode, try --uf-ho.";
       throw LogicException(ss.str());
     }
-    // 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 : expandDefinition : " << node << " to " << ret << std::endl;
+    Node ret = d_ho->expandDefinition(node);
+    if (ret != node)
+    {
+      Trace("uf-exp-def") << "TheoryUF::expandDefinition: higher-order: "
+                          << node << " to " << ret << std::endl;
       return ret;
     }
   }
@@ -391,18 +327,9 @@ bool TheoryUF::collectModelInfo(TheoryModel* m)
   }
 
   if( options::ufHo() ){
-    for( std::set<Node>::iterator it = termSet.begin(); it != termSet.end(); ++it ){
-      Node n = *it;
-      // for model-building with ufHo, we require that APPLY_UF is always
-      // expanded to HO_APPLY
-      if (!collectModelInfoHoTerm(n, m))
-      {
-        return false;
-      }
-    }
     // must add extensionality disequalities for all pairs of (non-disequal)
     // function equivalence classes.
-    if (checkExtensionality(m) != 0)
+    if (!d_ho->collectModelInfoHo(termSet, m))
     {
       return false;
     }
@@ -412,19 +339,6 @@ bool TheoryUF::collectModelInfo(TheoryModel* m)
   return true;
 }
 
-bool TheoryUF::collectModelInfoHoTerm(Node n, TheoryModel* m)
-{
-  if (n.getKind() == kind::APPLY_UF)
-  {
-    Node hn = TheoryUfRewriter::getHoApplyForApplyUf(n);
-    if (!m->assertEquality(n, hn, true))
-    {
-      return false;
-    }
-  }
-  return true;
-}
-
 void TheoryUF::presolve() {
   // TimerStat::CodeTimer codeTimer(d_presolveTimer);
 
@@ -749,236 +663,6 @@ void TheoryUF::eqNotifyDisequal(TNode t1, TNode t2, TNode reason) {
   }
 }
 
-Node TheoryUF::getExtensionalityDeq(TNode deq)
-{
-  Assert( deq.getKind()==kind::NOT && deq[0].getKind()==kind::EQUAL );
-  Assert( deq[0][0].getType().isFunction() );
-  std::map<Node, Node>::iterator it = d_extensionality_deq.find(deq);
-  if (it == d_extensionality_deq.end())
-  {
-    TypeNode tn = deq[0][0].getType();
-    std::vector<TypeNode> argTypes = tn.getArgTypes();
-    std::vector< Node > skolems;
-    NodeManager* nm = NodeManager::currentNM();
-    for (unsigned i = 0, nargs = argTypes.size(); i < nargs; i++)
-    {
-      Node k =
-          nm->mkSkolem("k", argTypes[i], "skolem created for extensionality.");
-      skolems.push_back( k );
-    }
-    Node t[2];
-    for (unsigned i = 0; i < 2; i++)
-    {
-      std::vector< Node > children;
-      Node curr = deq[0][i];
-      while( curr.getKind()==kind::HO_APPLY ){
-        children.push_back( curr[1] );
-        curr = curr[0];
-      }
-      children.push_back( curr );
-      std::reverse( children.begin(), children.end() );
-      children.insert( children.end(), skolems.begin(), skolems.end() );
-      t[i] = nm->mkNode(kind::APPLY_UF, children);
-    }
-    Node conc = t[0].eqNode( t[1] ).negate();
-    d_extensionality_deq[deq] = conc;
-    return conc;
-  }
-  return it->second;
-}
-
-unsigned TheoryUF::applyExtensionality(TNode deq)
-{
-  Assert(deq.getKind() == kind::NOT && deq[0].getKind() == kind::EQUAL);
-  Assert(deq[0][0].getType().isFunction());
-  // apply extensionality
-  if (d_extensionality.find(deq) == d_extensionality.end())
-  {
-    d_extensionality.insert(deq);
-    Node conc = getExtensionalityDeq(deq);
-    Node lem = NodeManager::currentNM()->mkNode( kind::OR, deq[0], conc );
-    Trace("uf-ho-lemma") << "uf-ho-lemma : extensionality : " << lem << std::endl;
-    d_out->lemma( lem );
-    return 1;
-  }
-  return 0;
-}
-
-unsigned TheoryUF::checkExtensionality(TheoryModel* m)
-{
-  unsigned num_lemmas = 0;
-  bool isCollectModel = (m != nullptr);
-  Trace("uf-ho") << "TheoryUF::checkExtensionality, collectModel="
-                 << isCollectModel << "..." << std::endl;
-  std::map< TypeNode, std::vector< Node > > func_eqcs;
-  eq::EqClassesIterator eqcs_i = eq::EqClassesIterator( &d_equalityEngine );
-  while( !eqcs_i.isFinished() ){
-    Node eqc = (*eqcs_i);
-    TypeNode tn = eqc.getType();
-    if( tn.isFunction() ){
-      // if during collect model, must have an infinite type
-      // if not during collect model, must have a finite type
-      if (tn.isInterpretedFinite() != isCollectModel)
-      {
-        func_eqcs[tn].push_back(eqc);
-        Trace("uf-ho-debug")
-            << "  func eqc : " << tn << " : " << eqc << std::endl;
-      }
-    }
-    ++eqcs_i;
-  }
-  
-  for( std::map< TypeNode, std::vector< Node > >::iterator itf = func_eqcs.begin(); 
-       itf != func_eqcs.end(); ++itf ){
-    for( unsigned j=0; j<itf->second.size(); j++ ){
-      for( unsigned k=(j+1); k<itf->second.size(); k++ ){ 
-        // if these equivalence classes are not explicitly disequal, do extensionality to ensure distinctness
-        if (!d_equalityEngine.areDisequal(
-                itf->second[j], itf->second[k], false))
-        {
-          Node deq = Rewriter::rewrite( itf->second[j].eqNode( itf->second[k] ).negate() );
-          // either add to model, or add lemma
-          if (isCollectModel)
-          {
-            // add extentionality disequality to the model
-            Node edeq = getExtensionalityDeq(deq);
-            Assert(edeq.getKind() == kind::NOT
-                   && edeq[0].getKind() == kind::EQUAL);
-            // introducing terms, must add required constraints, e.g. to
-            // force equalities between APPLY_UF and HO_APPLY terms
-            for (unsigned r = 0; r < 2; r++)
-            {
-              if (!collectModelInfoHoTerm(edeq[0][r], m))
-              {
-                return 1;
-              }
-            }
-            Trace("uf-ho-debug")
-                << "Add extensionality deq to model : " << edeq << std::endl;
-            if (!m->assertEquality(edeq[0][0], edeq[0][1], false))
-            {
-              return 1;
-            }
-          }
-          else
-          {
-            // apply extensionality lemma
-            num_lemmas += applyExtensionality(deq);
-          }
-        }
-      }   
-    }
-  }
-  return num_lemmas;
-}
-
-unsigned TheoryUF::applyAppCompletion( TNode n ) {
-  Assert( n.getKind()==kind::APPLY_UF );
-
-  //must expand into APPLY_HO version if not there already
-  Node ret = TheoryUfRewriter::getHoApplyForApplyUf( n );
-  if( !d_equalityEngine.hasTerm( ret ) || !d_equalityEngine.areEqual( ret, n ) ){
-    Node eq = ret.eqNode( n );
-    Trace("uf-ho-lemma") << "uf-ho-lemma : infer, by apply-expand : " << eq << std::endl;
-    d_equalityEngine.assertEquality(eq, true, d_true);
-    return 1;
-  }else{
-    Trace("uf-ho-debug") << "    ...already have " << ret << " == " << n << "." << std::endl;
-    return 0;
-  }
-}
-
-unsigned TheoryUF::checkAppCompletion() {
-  Trace("uf-ho") << "TheoryUF::checkApplyCompletion..." << std::endl;
-  // compute the operators that are relevant (those for which an HO_APPLY exist)
-  std::set< TNode > rlvOp;
-  eq::EqClassesIterator eqcs_i = eq::EqClassesIterator( &d_equalityEngine );
-  std::map< TNode, std::vector< Node > > apply_uf;
-  while( !eqcs_i.isFinished() ){
-    Node eqc = (*eqcs_i);
-    Trace("uf-ho-debug") << "  apply completion : visit eqc " << eqc << std::endl;
-    eq::EqClassIterator eqc_i = eq::EqClassIterator( eqc, &d_equalityEngine );
-    while( !eqc_i.isFinished() ){
-      Node n = *eqc_i;
-      if( n.getKind()==kind::APPLY_UF || n.getKind()==kind::HO_APPLY ){
-        int curr_sum = 0;
-        std::map< TNode, bool > curr_rops;
-        if( n.getKind()==kind::APPLY_UF ){
-          TNode rop = d_equalityEngine.getRepresentative( n.getOperator() );
-          if( rlvOp.find( rop )!=rlvOp.end() ){
-            // try if its operator is relevant
-            curr_sum = applyAppCompletion( n );
-            if( curr_sum>0 ){
-              return curr_sum;
-            }
-          }else{
-            // add to pending list
-            apply_uf[rop].push_back( n );
-          }
-          //arguments are also relevant operators  FIXME (github issue #1115)
-          for( unsigned k=0; k<n.getNumChildren(); k++ ){
-            if( n[k].getType().isFunction() ){
-              TNode rop = d_equalityEngine.getRepresentative( n[k] );
-              curr_rops[rop] = true;
-            }
-          }
-        }else{
-          Assert( n.getKind()==kind::HO_APPLY );
-          TNode rop = d_equalityEngine.getRepresentative( n[0] );
-          curr_rops[rop] = true;
-        }
-        for( std::map< TNode, bool >::iterator itc = curr_rops.begin(); itc != curr_rops.end(); ++itc ){
-          TNode rop = itc->first;
-          if( rlvOp.find( rop )==rlvOp.end() ){
-            rlvOp.insert( rop );
-            // now, try each pending APPLY_UF for this operator
-            std::map< TNode, std::vector< Node > >::iterator itu = apply_uf.find( rop );
-            if( itu!=apply_uf.end() ){
-              for( unsigned j=0; j<itu->second.size(); j++ ){
-                curr_sum = applyAppCompletion( itu->second[j] );
-                if( curr_sum>0 ){
-                  return curr_sum;
-                }
-              }
-            }
-          }
-        }
-      }
-      ++eqc_i;
-    }
-    ++eqcs_i;
-  }
-  return 0;
-}
-
-unsigned TheoryUF::checkHigherOrder() {
-  Trace("uf-ho") << "TheoryUF::checkHigherOrder..." << std::endl;
-
-  // infer new facts based on apply completion until fixed point 
-  unsigned num_facts;
-  do{
-    num_facts = checkAppCompletion();
-    if( d_conflict ){
-      Trace("uf-ho") << "...conflict during app-completion." << std::endl;
-      return 1;  
-    }
-  }while( num_facts>0 );
-
-  if( options::ufHoExt() ){
-    unsigned num_lemmas = 0;
-
-    num_lemmas = checkExtensionality();
-    if( num_lemmas>0 ){
-      Trace("uf-ho") << "...extensionality returned " << num_lemmas << " lemmas." << std::endl;
-      return num_lemmas;
-    }
-  }
-
-  Trace("uf-ho") << "...finished check higher order." << std::endl;
-
-  return 0;
-}
-
 } /* namespace CVC4::theory::uf */
 } /* namespace CVC4::theory */
 } /* namespace CVC4 */
diff --git a/src/theory/uf/theory_uf.h b/src/theory/uf/theory_uf.h
index 72dc04b10..df1cc232b 100644
--- a/src/theory/uf/theory_uf.h
+++ b/src/theory/uf/theory_uf.h
@@ -20,6 +20,7 @@
 #ifndef CVC4__THEORY__UF__THEORY_UF_H
 #define CVC4__THEORY__UF__THEORY_UF_H
 
+#include "context/cdhashmap.h"
 #include "context/cdhashset.h"
 #include "context/cdo.h"
 #include "expr/node.h"
@@ -34,6 +35,7 @@ namespace uf {
 
 class UfTermDb;
 class StrongSolverTheoryUF;
+class HoExtension;
 
 class TheoryUF : public Theory {
 
@@ -130,14 +132,8 @@ private:
   /** The conflict node */
   Node d_conflictNode;
 
-  /** extensionality has been applied to these disequalities */
-  NodeSet d_extensionality;
-
-  /** cache of getExtensionalityDeq below */
-  std::map<Node, Node> d_extensionality_deq;
-
-  /** map from non-standard operators to their skolems */
-  NodeNodeMap d_uf_std_skolem;
+  /** the higher-order solver extension */
+  HoExtension* d_ho;
 
   /** node for true */
   Node d_true;
@@ -180,85 +176,7 @@ private:
   /** called when two equivalence classes are made disequal */
   void eqNotifyDisequal(TNode t1, TNode t2, TNode reason);
 
- private:  // for higher-order
-  /** get extensionality disequality
-   *
-   * Given disequality deq f != g, this returns the disequality:
-   *   (f k) != (g k) for fresh constant(s) k.
-   */
-  Node getExtensionalityDeq(TNode deq);
-
-  /** applyExtensionality
-   *
-   * Given disequality deq f != g, if not already cached, this sends a lemma of
-   * the form:
-   *   f = g V (f k) != (g k) for fresh constant k.
-   * on the output channel. This is an "extensionality lemma".
-   * Return value is the number of lemmas of this form sent on the output
-   * channel.
-   */
-  unsigned applyExtensionality(TNode deq);
-
-  /**
-   * Check whether extensionality should be applied for any pair of terms in the
-   * equality engine.
-   *
-   * If we pass a null model m to this function, then we add extensionality
-   * lemmas to the output channel and return the total number of lemmas added.
-   * We only add lemmas for functions whose type is finite, since pairs of
-   * functions whose types are infinite can be made disequal in a model by
-   * witnessing a point they are disequal.
-   *
-   * If we pass a non-null model m to this function, then we add disequalities
-   * that correspond to the conclusion of extensionality lemmas to the model's
-   * equality engine. We return 0 if the equality engine of m is consistent
-   * after this call, and 1 otherwise. We only add disequalities for functions
-   * whose type is infinite, since our decision procedure guarantees that
-   * extensionality lemmas are added for all pairs of functions whose types are
-   * finite.
-   */
-  unsigned checkExtensionality(TheoryModel* m = nullptr);
-
-  /** applyAppCompletion
-   * This infers a correspondence between APPLY_UF and HO_APPLY
-   * versions of terms for higher-order.
-   * Given APPLY_UF node e.g. (f a b c), this adds the equality to its
-   * HO_APPLY equivalent:
-   *   (f a b c) == (@ (@ (@ f a) b) c)
-   * to equality engine, if not already equal.
-   * Return value is the number of equalities added.
-   */
-  unsigned applyAppCompletion(TNode n);
-
-  /** check whether app-completion should be applied for any
-   * pair of terms in the equality engine.
-   */
-  unsigned checkAppCompletion();
-
-  /** check higher order
-   * This is called at full effort and infers facts and sends lemmas
-   * based on higher-order reasoning (specifically, extentionality and
-   * app completion above). It returns the number of lemmas plus facts
-   * added to the equality engine.
-  */
-  unsigned checkHigherOrder();
-
-  /** collect model info for higher-order term
-   *
-   * This adds required constraints to m for term n. In particular, if n is
-   * an APPLY_UF term, we add its HO_APPLY equivalent in this call. We return
-   * true if the model m is consistent after this call.
-   */
-  bool collectModelInfoHoTerm(Node n, TheoryModel* m);
-
-  /** get apply uf for ho apply
-   * This returns the APPLY_UF equivalent for the HO_APPLY term node, where
-   * node has non-functional return type (that is, it corresponds to a fully
-   * applied function term).
-   * This call may introduce a skolem for the head operator and send out a lemma
-   * specifying the definition.
-  */
-  Node getApplyUfForHoApply(Node node);
+ private:
   /** get the operator for this node (node should be either APPLY_UF or
    * HO_APPLY)
    */
@@ -303,7 +221,10 @@ private:
   StrongSolverTheoryUF* getStrongSolver() {
     return d_thss;
   }
-private:
+  /** are we in conflict? */
+  bool inConflict() const { return d_conflict; }
+
+ private:
   bool areCareDisequal(TNode x, TNode y);
   void addCarePairs(TNodeTrie* t1,
                     TNodeTrie* t2,