Merge pull request #656 from glu-lang/sema-splitting

[Sema] Feature: CS Splitting

Author: Snowy1803

include/Sema/Constraint.hpp

@@ -455,6 +455,15 @@ class Constraint {
     /// @return True if the constraint is disabled, false otherwise.
     bool isDisabled() const { return _isDisabled; }
 
+    /// @brief Disable this constraint.
+    void disable() { _isDisabled = 1; }
+    /// @brief Enable this constraint.
+    void enable() { _isDisabled = 0; }
+
+    /// @brief Sets the enabled state of this constraint.
+    /// @param enabled True to enable the constraint, false to disable it.
+    void setEnabled(bool enabled) { _isDisabled = !enabled; }
+
     /// @brief Print this constraint to the output stream.
     void print() const;
 };

include/Sema/ConstraintSystem.hpp

@@ -6,6 +6,7 @@
 #include "ScopeTable.hpp"
 
 #include <llvm/ADT/DenseMap.h>
+#include <llvm/ADT/DenseSet.h>
 #include <llvm/ADT/ilist.h>
 
 namespace glu::sema {
@@ -58,6 +59,10 @@ struct SystemState {
     /// @return A deep copy of the current state.
     SystemState clone() const { return *this; }
 
+    /// @brief Merges this state into another, combining bindings and choices.
+    /// @param other The target state to merge into.
+    void mergeInto(SystemState &other) const;
+
     /// @brief Calculates the score of the current state based on implicit
     /// conversions.
     /// @return The score representing the number of implicit conversions.
@@ -214,8 +219,16 @@ class ConstraintSystem {
     /// extracted.
     void mapImplicitConversions(Solution *solution);
 
-    /// @brief Solves all constraints and applies type mappings to the specified
-    /// expressions.
+    /// @brief Solves all enabled constraints within this local constraint
+    /// system and returns the solution result through the provided parameter.
+    /// This method is called within solveConstraints, after simplification and
+    /// splitting of the constraint system, and before mapping types back to the
+    /// AST.
+    /// @param result The solution result to populate with found solutions.
+    /// @return True if a solution was found, false otherwise.
+    bool solveLocalConstraints(SolutionResult &result);
+
+    /// @brief Solves all constraints and applies mappings.
     ///
     /// This method combines constraint solving with type mapping for
     /// expressions. For module expressions (part of the AST tree), type
@@ -424,6 +437,11 @@ void printConstraints(
     ConstraintSystem &system, llvm::raw_ostream &os = llvm::outs()
 );
 
+void collectTypeVariables(
+    Constraint *constraint,
+    llvm::DenseSet<glu::types::TypeVariableTy *> &typeVars
+);
+
 } // namespace glu::sema
 
 #endif // GLU_SEMA_CONSTRAINT_SYSTEM_HPP

lib/Sema/CMakeLists.txt

@@ -25,6 +25,7 @@ target_sources(Sema
         ConstraintSystem/OccursCheckVisitor.cpp
         ConstraintSystem/Solver.cpp
         ConstraintSystem/SubstitutionMapper.cpp
+        ConstraintSystem/TypeVariableCollector.cpp
         ConstraintSystem/TypeVariableTyMapper.cpp
         ConstraintSystem/UnificationVisitor.cpp
         ModuleWalker.cpp

lib/Sema/ConstraintSystem/ConstraintSystem.cpp

@@ -145,20 +145,16 @@ void ConstraintSystem::mapImplicitConversions(Solution *solution)
     }
 }
 
-bool ConstraintSystem::solveConstraints()
+bool ConstraintSystem::solveLocalConstraints(SolutionResult &result)
 {
     /// The initial system state used to begin constraint solving.
     std::vector<SystemState> worklist;
     worklist.emplace_back(_context); // Start from an empty state
 
-    SolutionResult result; // Local solution result
-
     while (!worklist.empty()) {
         SystemState current = std::move(worklist.back());
         worklist.pop_back();
 
-        bool failed = false;
-
         /// Apply non-defaultable constraints first
         for (Constraint *constraint : _constraints) {
             // Skip disabled constraints
@@ -176,15 +172,11 @@ bool ConstraintSystem::solveConstraints()
             ConstraintResult result = apply(constraint, current, worklist);
             markConstraint(result, constraint);
             if (result == ConstraintResult::Failed) {
-                failed = true;
-                break;
+                goto failed;
             }
             // Continue if Satisfied or Applied
         }
 
-        if (failed)
-            continue;
-
         for (Constraint *constraint : _constraints) {
             if (constraint->isDisabled())
                 continue;
@@ -194,8 +186,7 @@ bool ConstraintSystem::solveConstraints()
             ConstraintResult result = apply(constraint, current, worklist);
             markConstraint(result, constraint);
             if (result == ConstraintResult::Failed) {
-                failed = true;
-                break;
+                goto failed;
             }
         }
 
@@ -214,8 +205,7 @@ bool ConstraintSystem::solveConstraints()
             ConstraintResult result = apply(constraint, current, worklist);
             markConstraint(result, constraint);
             if (result == ConstraintResult::Failed) {
-                failed = true;
-                break;
+                goto failed;
             }
             // Continue if Satisfied or Applied
         }
@@ -235,17 +225,16 @@ bool ConstraintSystem::solveConstraints()
             ConstraintResult result = apply(constraint, current, worklist);
             markConstraint(result, constraint);
             if (result == ConstraintResult::Failed) {
-                failed = true;
-                break;
+                goto failed;
             }
             // Continue if Satisfied or Applied
         }
 
-        if (failed)
-            continue;
-
         /// All constraints are satisfied -- record the solution.
         result.tryAddSolution(current);
+
+    failed:
+        continue;
     }
 
     if (result.isAmbiguous()) {
@@ -259,9 +248,68 @@ bool ConstraintSystem::solveConstraints()
         reportNoSolutionError();
         return false;
     }
-    mapTypeVariables(solution);
-    mapOverloadChoices(solution);
-    mapImplicitConversions(solution);
+    return true;
+}
+
+bool ConstraintSystem::solveConstraints()
+{
+    // color the constraints based on which type variables they contain
+    // map of color => set of used type variables
+    std::vector<llvm::DenseSet<glu::types::TypeVariableTy *>> colors;
+    // map of color => set of constraints
+    std::vector<llvm::DenseSet<Constraint *>> colorConstraints;
+    for (auto *constraint : _constraints) {
+        llvm::DenseSet<glu::types::TypeVariableTy *> typeVars;
+        collectTypeVariables(constraint, typeVars);
+        colors.push_back(std::move(typeVars));
+        colorConstraints.push_back({ constraint });
+    }
+    // merge colors that share type variables
+    bool changed;
+    do {
+        changed = false;
+        for (std::size_t i = 0; i < colors.size(); ++i) {
+            for (std::size_t j = i + 1; j < colors.size(); ++j) {
+                llvm::DenseSet<glu::types::TypeVariableTy *> intersection;
+                for (auto *typeVar : colors[j]) {
+                    if (colors[i].count(typeVar)) {
+                        intersection.insert(typeVar);
+                    }
+                }
+                if (!intersection.empty()) {
+                    // merge j into i
+                    for (auto *typeVar : colors[j]) {
+                        colors[i].insert(typeVar);
+                    }
+                    for (auto *constraint : colorConstraints[j]) {
+                        colorConstraints[i].insert(constraint);
+                    }
+                    changed = true;
+                    // Erase the merged color
+                    colors.erase(colors.begin() + j);
+                    colorConstraints.erase(colorConstraints.begin() + j);
+                    --j;
+                }
+            }
+        }
+    } while (changed);
+    // solve each color separately
+    SystemState finalSolution(_context);
+    for (std::size_t i = 0; i < colors.size(); ++i) {
+        // Disable all constraints not in this color
+        for (auto *constraint : _constraints) {
+            constraint->setEnabled(colorConstraints[i].count(constraint));
+        }
+        SolutionResult result;
+        if (!solveLocalConstraints(result)) {
+            return false;
+        }
+        result.getBestSolution()->mergeInto(finalSolution);
+    }
+
+    mapTypeVariables(&finalSolution);
+    mapOverloadChoices(&finalSolution);
+    mapImplicitConversions(&finalSolution);
     return true;
 }
 

lib/Sema/ConstraintSystem/Solver.cpp

@@ -148,4 +148,22 @@ void SolutionResult::tryAddSolution(SystemState const &s)
     }
 }
 
+void SystemState::mergeInto(SystemState &other) const
+{
+    // Merge type bindings
+    for (auto const &[var, type] : typeBindings) {
+        other.typeBindings[var] = type;
+    }
+
+    // Merge overload choices
+    for (auto const &[expr, decl] : overloadChoices) {
+        other.overloadChoices[expr] = decl;
+    }
+
+    // Merge implicit conversions
+    for (auto const &[expr, targetType] : implicitConversions) {
+        other.implicitConversions[expr] = targetType;
+    }
+}
+
 } // namespace glu::sema

lib/Sema/ConstraintSystem/TypeVariableCollector.cpp

@@ -0,0 +1,73 @@
+#include "AST/Types.hpp"
+#include "Constraint.hpp"
+
+#include <llvm/ADT/DenseSet.h>
+
+namespace glu::sema {
+
+class TypeVariableCollector
+    : public glu::types::TypeVisitor<TypeVariableCollector, void> {
+protected:
+    llvm::DenseSet<glu::types::TypeVariableTy *> &_typeVariables;
+
+public:
+    TypeVariableCollector(
+        llvm::DenseSet<glu::types::TypeVariableTy *> &typeVariables
+    )
+        : _typeVariables(typeVariables)
+    {
+    }
+
+    void visitTypeBase([[maybe_unused]] glu::types::TypeBase *type) { }
+
+    void visitFunctionTy(glu::types::FunctionTy *type)
+    {
+        visit(type->getReturnType());
+        for (glu::types::TypeBase *paramType : type->getParameters())
+            visit(paramType);
+    }
+
+    void visitPointerTy(types::PointerTy *type) { visit(type->getPointee()); }
+
+    void visitTypeAliasTy(types::TypeAliasTy *type)
+    {
+        visit(type->getWrappedType());
+    }
+
+    void visitStaticArrayTy(types::StaticArrayTy *type)
+    {
+        visit(type->getDataType());
+    }
+
+    void visitDynamicArrayTy(types::DynamicArrayTy *type)
+    {
+        visit(type->getDataType());
+    }
+
+    void visitTypeVariableTy(glu::types::TypeVariableTy *type)
+    {
+        _typeVariables.insert(type);
+    }
+};
+
+void collectTypeVariables(
+    Constraint *constraint,
+    llvm::DenseSet<glu::types::TypeVariableTy *> &typeVars
+)
+{
+    if (constraint->getKind() == ConstraintKind::Disjunction
+        || constraint->getKind() == ConstraintKind::Conjunction) {
+        for (Constraint *nested : constraint->getNestedConstraints()) {
+            collectTypeVariables(nested, typeVars);
+        }
+        return;
+    }
+    TypeVariableCollector collector(typeVars);
+    collector.visit(constraint->getFirstType());
+    if (!constraint->isTypePropertyConstraint()
+        && constraint->getKind() != ConstraintKind::BindOverload
+        && constraint->getKind() != ConstraintKind::StructInitialiser)
+        collector.visit(constraint->getSecondType());
+}
+
+} // namespace glu::sema