Sum Pyramids logic programming solution

sum_pyramids.rb

require "mini_kanren"

def suc(nat)
  [:suc, nat]
end

def int_to_nat(i)
  i == 0 ? :zero : suc(int_to_nat(i - 1))
end

def nat_to_int(n)
  case n
  when String
    puts "Undecided natural number found, the solution is divergent"
    exit 1
  when :zero
    0
  else
    1 + nat_to_int(n[1])
  end
end

def plus(a, b, sum)
  ->(s) {
    conde(
      all(eq(a, int_to_nat(0)), eq(b, sum)),
      fresh { |a_, sum_| all(eq(suc(a_), a), eq(suc(sum_), sum), plus(a_, b, sum_)) }
    ).call(s)
  }
end

_, *input = $<.map { |l| l.split.map { _1 == "?" ? nil : _1.to_i } }

solutions = MiniKanren.exec do
  goals = []

  pyramid = fresh
  pyramid_slices = fresh(input.size)
  goals << eq(pyramid, pyramid_slices)

  var_rows =
    input.zip(pyramid_slices).map do |input_row, slice|
      var_row = fresh(input_row.size)
      goals << eq(slice, var_row)
      input_row.zip(var_row).each { |int, var| goals << eq(var, int_to_nat(int)) if int }
      var_row
    end

  var_rows.each_cons(2).each do |row, next_row|
    row.size.times { |i| goals << plus(next_row[i], next_row[i + 1], row[i]) }
  end

  run(pyramid, *goals)
end

solutions.each do |solution|
  puts solution.map { |row| row.map { nat_to_int(_1) }.join(" ") }
end