2 files changed, 187 insertions, 0 deletions

diff --git a/src/infer/substitution.cc b/src/infer/substitution.cc
new file mode 100644
index 0000000..c570676
--- /dev/null
+++ b/src/infer/substitution.cc
@@ -0,0 +1,119 @@
+/*
+ * richard-atp-engine: A saturation-based automated theorem prover for first-order logic.
+ * Copyright (C) 2026 Richard Qin
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+#include "atp/infer/substitution.h"
+#include "atp/core/term.h"
+#include "atp/core/term_bank.h"
+#include "atp/core/types.h"
+
+namespace atp {
+
+bool Substitution::bind(TermId var, TermId term) {
+    auto [it, inserted] = bindings_.emplace(var, term);
+    if (!inserted) {
+        // Already bound — succeed only if bound to the same term
+        return it->second == term;
+    }
+    return true;
+}
+
+TermId Substitution::lookup(TermId var) const {
+    auto it = bindings_.find(var);
+    if (it != bindings_.end()) {
+        return it->second;
+    }
+    return kInvalidId;
+}
+
+TermId Substitution::resolve(TermId id, const TermBank& bank) const {
+    // Walk the chain: X → Y → Z → f(a)
+    // Stop when we hit a non-variable or an unbound variable.
+    while (bank.isVariable(id)) {
+        auto it = bindings_.find(id);
+        if (it == bindings_.end()) {
+            break;  // Unbound variable — stop here
+        }
+        id = it->second;
+    }
+    return id;
+}
+
+TermId Substitution::apply(TermId term, TermBank& bank) const {
+    // First, resolve through variable chain
+    term = resolve(term, bank);
+
+    if (bank.isVariable(term)) {
+        // Free variable (unbound) — leave as-is
+        return term;
+    }
+
+    // IMPORTANT: Copy symbol_id and args BEFORE recursing.
+    // Recursive apply() calls may trigger bank.makeTerm() → realloc → dangling ref.
+    const Term& t = bank.getTerm(term);
+    SymbolId sym = t.symbol_id;
+    std::vector<TermId> old_args(t.args);  // copy!
+
+    if (old_args.empty()) {
+        // Constant or 0-arity — no substitution needed
+        return term;
+    }
+
+    // Recursively apply to all arguments
+    std::vector<TermId> new_args;
+    new_args.reserve(old_args.size());
+    bool changed = false;
+    for (TermId arg : old_args) {
+        TermId new_arg = apply(arg, bank);
+        new_args.push_back(new_arg);
+        if (new_arg != arg) {
+            changed = true;
+        }
+    }
+
+    if (!changed) {
+        return term;
+    }
+
+    return bank.makeTerm(sym, new_args);
+}
+
+bool Substitution::occursIn(TermId var, TermId term, const TermBank& bank) const {
+    term = resolve(term, bank);
+
+    if (var == term) {
+        return true;
+    }
+
+    if (bank.isVariable(term)) {
+        return false;  // Different unbound variable
+    }
+
+    const Term& t = bank.getTerm(term);
+    for (TermId arg : t.args) {
+        if (occursIn(var, arg, bank)) {
+            return true;
+        }
+    }
+    return false;
+}
+
+void Substitution::clear() {
+    bindings_.clear();
+}
+
+}  // namespace atp
\ No newline at end of file
diff --git a/src/infer/unification.cc b/src/infer/unification.cc
index b62ffd0..be8a025 100644
--- a/src/infer/unification.cc
+++ b/src/infer/unification.cc
@@ -17,3 +17,71 @@
  */
 
 
+#include "atp/infer/unification.h"
+#include "atp/core/types.h"
+#include "atp/infer/substitution.h"
+
+namespace atp {
+
+    bool unify(const TermBank& bank, TermId t1, TermId t2, Substitution& subst,
+           const UnificationConfig& config) {
+        t1 = subst.resolve(t1, bank);
+        t2 = subst.resolve(t2, bank);
+
+        if (t1 == t2) {
+            return true;
+        }
+
+        if (bank.isVariable(t1)) {
+            if (config.enable_occurs_check && subst.occursIn(t1, t2, bank)) {
+                return false;
+            }
+            subst.bind(t1, t2);
+            return true;
+        }
+
+        if (bank.isVariable(t2)) {
+            if (config.enable_occurs_check && subst.occursIn(t2, t1, bank)) {
+                return false;
+            }
+            subst.bind(t2, t1);
+            return true;
+        }
+
+        const Term& term1 = bank.getTerm(t1);
+        const Term& term2 = bank.getTerm(t2);
+
+        if (term1.symbol_id != term2.symbol_id || term1.arity() != term2.arity()) {
+            return false;
+        }
+
+        for (size_t i = 0; i < term1.arity(); ++i) {
+            if (!unify(bank, term1.args[i], term2.args[i], subst, config)) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    TermId applySubstitution(const Substitution& subst, TermId term, TermBank& bank) {
+        term = subst.resolve(term, bank);
+        if (bank.isVariable(term)) {
+            return term;
+        }
+
+        // IMPORTANT: Copy args and symbol_id BEFORE recursing.
+        // Recursive calls may call bank.makeTerm() → realloc terms_ → invalidate refs.
+        const Term& t = bank.getTerm(term);
+        SymbolId sym = t.symbol_id;
+        std::vector<TermId> old_args(t.args);  // copy!
+
+        std::vector<TermId> new_args;
+        new_args.reserve(old_args.size());
+        for (TermId arg : old_args) {
+            new_args.push_back(applySubstitution(subst, arg, bank));
+        }
+        return bank.makeTerm(sym, new_args);
+    }
+
+}  // namespace atp
\ No newline at end of file