Constraint Solver WIP

constraint.rb

class Var < Struct.new(:name, :domain)
  def self.gen(domains)
    domain = domains.first.product(*domains[1..-1])
    Var.new("__gen__#{domain.object_id}", domain)
  end

  def empty?
    domain.empty?
  end

  def assignments
    domain.map { |val| assign(val) }
  end

  def assign(value)
    Var.new(name, [value])
  end

  def assigned?
    domain.size == 1
  end

  def gen?
    name =~ /__gen__/
  end

  def propogate(constraints, env)
    constraints.reduce(env) do |env, constraint|
      if constraint.applies?(name)
        constraint.apply(env)
      else
        env
      end
    end
  end
end

class UnaryConstraint < Struct.new(:name, :condition)
  def applies?(var_name)
    var_name == name
  end

  def apply(env)
    var = env[name]
    new_domain = var.domain.select { |val| condition.call(val) }
    env.merge({ name => Var.new(name, new_domain) })
  end
end

class BinaryConstraint < Struct.new(:name_1, :name_2, :condition)
  def self.from_multi(names, env, condition)
    domains = names.map { |name| env[name].domain }
    var = Var.gen(domains)
    constraints = names.map.with_index do |name, i|
      constraint = Proc.new { |x, combo| x == combo[i] && condition.call(*combo) }
      BinaryConstraint.new(name, var.name, constraint)
    end
    [var, constraints]
  end

  def applies?(var_name)
    var_name == name_1 || var_name == name_2
  end

  def apply(env)
    var_1 = env[name_1]
    var_2 = env[name_2]
    new_domain_1 = var_1.domain.select do |v1|
      var_2.domain.any? do |v2|
        condition.call(v1, v2)
      end
    end
    new_domain_2 = var_2.domain.select do |v2|
      var_1.domain.any? do |v1|
        condition.call(v1, v2)
      end
    end
    env.merge({
      name_1 => Var.new(name_1, new_domain_1),
      name_2 => Var.new(name_2, new_domain_2)
    })
  end
end

class Solution < Struct.new(:assignments)
  NONE = Object.new

  def self.none
    Solution.new(NONE)
  end

  def to_hash
    assignments.reduce({}) do |h, var|
      h[var.name] = var.domain.first unless var.gen?
      h
    end
  end

  def found?
    assignments != NONE
  end
end

class Problem < Struct.new(:env, :constraints)
  def var(name, domain)
    env[name] = Var.new(name, domain)
  end

  def vars
    env.values
  end

  def constrain(*names, &condition)
    case names.size
    when 1
      constraints << UnaryConstraint.new(names.first, condition)
    when 2
      constraints << BinaryConstraint.new(names.first, names.last, condition)
    else
      new_var, new_constraints = BinaryConstraint.from_multi(names, env, condition)
      env[new_var.name] = new_var
      new_constraints.each { |c| constraints << c }
    end
  end

  def solve
    new_env = vars.reduce(env) { |env, var| var.propogate(constraints, env) }
    constrained = new_env.values
    case
    when constrained.all?(&:assigned?)
      Solution.new(vars)
    when constrained.any?(&:empty?)
      Solution.none
    else
      solve_tree(constrained, new_env)
    end
  end

  def solve_tree(constrained_vars, new_env)
    constrained_vars
      .sort_by { |var| var.domain.size }
      .find { |var| !var.assigned? }
      .assignments
      .map { |var| Problem.new(new_env.merge({ var.name => var }), constraints) }
      .map(&:solve)
      .flatten
      .select(&:found?)
  end
end

# Simple
prob = Problem.new({}, [])
prob.var(:a, (0..4).to_a)
prob.var(:b, (0..4).to_a)
prob.var(:c, (0..4).to_a)
prob.constrain(:a, :b) { |a, b| a < b }
prob.constrain(:a, :b, :c) { |a, b, c| a + b == c }
p prob.solve.map(&:to_hash)

# Magic square
prob = Problem.new({}, [])
squares = (0..15).to_a
squares.each do |square|
  prob.var(square, (1..16).to_a)
end
sum_constraint = ->(*line) { line.reduce(:+) == 34 }
prob.constrain(squares) { |*squares| squares.size == squares.uniq.size }
prob.constrain([0, 5, 10, 15], &sum_constraint)
prob.constrain([3, 6, 9, 12], &sum_constraint)
(0..3).each do |i|
  row = (0..3).map { |j| i * 4 + j }
  prob.constrain(row, &sum_constraint)
end
(0..3).each do |i|
  col = (0..3).map { |j| i + 4 * j }
  prob.constrain(col, &sum_constraint)
end
p prob.solve