diff --git a/src/ast/rewriter/finite_set_axioms.cpp b/src/ast/rewriter/finite_set_axioms.cpp
index 411fc6d8a..dad7dd162 100644
--- a/src/ast/rewriter/finite_set_axioms.cpp
+++ b/src/ast/rewriter/finite_set_axioms.cpp
@@ -32,6 +32,27 @@ std::ostream& operator<<(std::ostream& out, theory_axiom const& ax) {
     return out;
 }
 
+void finite_set_axioms::add_unit(char const *name, expr *p1, expr *unit) {
+    theory_axiom *ax = alloc(theory_axiom, m, name, p1);
+    ax->clause.push_back(unit);
+    m_add_clause(ax);
+}
+
+void finite_set_axioms::add_binary(char const *name, expr *p1, expr *p2, expr *f1, expr *f2) {
+    theory_axiom *ax = alloc(theory_axiom, m, name, p1, p2);
+    ax->clause.push_back(f1);
+    ax->clause.push_back(f2);
+    m_add_clause(ax);
+}
+
+void finite_set_axioms::add_ternary(char const *name, expr *p1, expr *p2, expr *f1, expr *f2, expr *f3) {
+    theory_axiom *ax = alloc(theory_axiom, m, name, p1, p2);
+    ax->clause.push_back(f1);
+    ax->clause.push_back(f2);
+    ax->clause.push_back(f3);
+    m_add_clause(ax);
+}
+
 // a ~ set.empty => not (x in a)
 // x is an element, generate axiom that x is not in any empty set of x's type
 void finite_set_axioms::in_empty_axiom(expr *x) {
@@ -205,6 +226,7 @@ void finite_set_axioms::in_range_axiom(expr* r) {
 
 // a := set.map(f, b)
 // (x in a) <=> set.map_inverse(f, x, b) in b
+// 
 void finite_set_axioms::in_map_axiom(expr *x, expr *a) {
     expr *f = nullptr, *b = nullptr;
     sort *elem_sort = nullptr;
@@ -270,27 +292,6 @@ void finite_set_axioms::in_filter_axiom(expr *x, expr *a) {
     add_ternary("in-filter", x, a, m.mk_not(x_in_b), npx, x_in_a);
 }
 
-void finite_set_axioms::add_unit(char const* name, expr* p1, expr* unit) {
-    theory_axiom *ax = alloc(theory_axiom, m, name, p1);
-    ax->clause.push_back(unit);
-    m_add_clause(ax);
-}
-
-void finite_set_axioms::add_binary(char const* name, expr* p1, expr* p2, expr* f1, expr* f2) {
-    theory_axiom *ax = alloc(theory_axiom, m, name, p1, p2);
-    ax->clause.push_back(f1);
-    ax->clause.push_back(f2);
-    m_add_clause(ax);
-}
-
-void finite_set_axioms::add_ternary(char const *name, expr *p1, expr *p2, expr *f1, expr *f2, expr *f3) {
-    theory_axiom *ax = alloc(theory_axiom, m, name, p1, p2);
-    ax->clause.push_back(f1);
-    ax->clause.push_back(f2);
-    ax->clause.push_back(f3);
-    m_add_clause(ax);
-}
-
 // Auxiliary algebraic axioms to ease reasoning about set.size
 // The axioms are not required for completenss for the base fragment
 // as they are handled by creating semi-linear sets.
diff --git a/src/smt/theory_finite_set.cpp b/src/smt/theory_finite_set.cpp
index fd4715470..1bb6a9c90 100644
--- a/src/smt/theory_finite_set.cpp
+++ b/src/smt/theory_finite_set.cpp
@@ -631,36 +631,33 @@ namespace smt {
     void theory_finite_set::add_membership_axioms(expr *elem, expr *set) {
         TRACE(finite_set, tout << "add_membership_axioms: " << mk_pp(elem, m) << " in " << mk_pp(set, m) << "\n";);
 
-            // Instantiate appropriate axiom based on set structure
+        // Instantiate appropriate axiom based on set structure
         if (!is_new_axiom(elem, set))
             ;
-        else if (u.is_empty(set)) {
-            m_axioms.in_empty_axiom(elem);
-        }
-        else if (u.is_singleton(set)) {
-            m_axioms.in_singleton_axiom(elem, set);
-        }
-        else if (u.is_union(set)) {
-            m_axioms.in_union_axiom(elem, set);
-        }
-        else if (u.is_intersect(set)) {
-            m_axioms.in_intersect_axiom(elem, set);
-        }
-        else if (u.is_difference(set)) {
-            m_axioms.in_difference_axiom(elem, set);
-        }
-        else if (u.is_range(set)) {
-            m_axioms.in_range_axiom(elem, set);
-        }
+        else if (u.is_empty(set))
+            m_axioms.in_empty_axiom(elem);        
+        else if (u.is_singleton(set)) 
+            m_axioms.in_singleton_axiom(elem, set);        
+        else if (u.is_union(set)) 
+            m_axioms.in_union_axiom(elem, set);        
+        else if (u.is_intersect(set)) 
+            m_axioms.in_intersect_axiom(elem, set);        
+        else if (u.is_difference(set)) 
+            m_axioms.in_difference_axiom(elem, set);        
+        else if (u.is_range(set))
+            m_axioms.in_range_axiom(elem, set);        
         else if (u.is_map(set)) {
-            // TODO type of elem could be from the pre-image
+            // sort *elem_sort = u.finte_set_elem_sort(set->get_sort());
+            
+            // set.map_inverse can loop. need to check instance generation.
             m_axioms.in_map_axiom(elem, set);
+            
+            // this can also loop:
             m_axioms.in_map_image_axiom(elem, set);
         }
-        else if (u.is_filter(set)) {            
+        else if (u.is_filter(set)) {
             m_axioms.in_filter_axiom(elem, set);
         }
-        TRACE(finite_set, tout << "after add_membership_axioms: " << mk_pp(elem, m) << " in " << mk_pp(set, m) << "\n";);
     }
 
     void theory_finite_set::add_clause(theory_axiom* ax) {