Skip to content

Instantly share code, notes, and snippets.

@davidbalbert
Last active April 4, 2024 19:47
Show Gist options
  • Save davidbalbert/dd03c945752c78f7f0a5fbac45a17c03 to your computer and use it in GitHub Desktop.
Save davidbalbert/dd03c945752c78f7f0a5fbac45a17c03 to your computer and use it in GitHub Desktop.
miniKanren
// miniKanren.swift, written by David Albert, and released into the public domain.
//
// An implementation of the miniKanren relational programming language. Somewhat
// reductively, relational programs can be run both forwards and backwards. In
// addition to asking "what does this function return when passed these arguments,"
// you can also ask "what arguments to this function return the following value?"
// This lets you ask questions like like "what pairs of numbers sum to 12," which
// if you allow for negative numbers has an infinite number of solutions.
//
// Some mind bending miniKanren programs that others have written include: a
// Scheme interpreter that can be run backwards – i.e. "give me all the Scheme
// programs that return the list (1 2 3)" – and a MIPS assembler that's
// simultaneously a disassembler. After all, a disassembler is just an assembler
// run in reverse.
//
// Implemented while reading The Reasoned Schemer by Daniel Friedman, et al. For
// example usage, see the end of the file.
// µKanren - the smallest relational programming language.
// A logic variable. Basically just a unique index. Two Vars with the same index
// are the same, and thus must be bound to the same value. It's possible for a
// variable to be unbound, which means it can take on any value.
struct Var: Equatable, Hashable {
let i: Int
init(_ i: Int) {
self.i = i
}
}
// A value or variable.
enum Term: Equatable {
case none
case int(Int)
case double(Double)
case string(String)
case bool(Bool)
indirect case pair(Term, Term) // a cons cell
case `var`(Var)
// A string representation of an unbound variable. Only present after the
// outermost expression in a program has been evaluated.
case sym(String)
}
extension Term: CustomStringConvertible {
var description: String {
switch self {
case .int(let i):
return "\(i)"
case .double(let d):
return "\(d)"
case .string(let s):
return "\"\(s)\""
case .bool(let b):
return "\(b)"
case .var(let v):
return "\(v)"
case .sym(let s):
return s
case .none:
return "()"
case .pair(_, _):
var elements = [String]()
var l = self
while case .pair(let first, let rest) = l {
elements.append("\(first)")
l = rest
}
// improper list
if l != Term.none {
elements.append(". \(l)")
}
return "(\(elements.joined(separator: " ")))"
}
}
}
extension Term: ExpressibleByIntegerLiteral {
init(integerLiteral: Int) {
self = .int(integerLiteral)
}
}
extension Term: ExpressibleByFloatLiteral {
init(floatLiteral: Double) {
self = .double(floatLiteral)
}
}
extension Term: ExpressibleByStringLiteral {
init(stringLiteral: String) {
self = .string(stringLiteral)
}
}
extension Term: ExpressibleByBooleanLiteral {
init(booleanLiteral: Bool) {
self = .bool(booleanLiteral)
}
}
extension Term: ExpressibleByArrayLiteral {
init(arrayLiteral elements: Term...) {
var l: Term = .none
for e in elements.reversed() {
l = .pair(e, l)
}
self = l
}
}
extension Term: Sequence {
struct Iterator: IteratorProtocol {
var term: Term
mutating func next() -> Term? {
switch term {
case .pair(let value, let next):
defer { term = next }
return value
default:
return nil
}
}
}
public func makeIterator() -> Iterator {
return Iterator(term: self)
}
}
// A dictionary binding logic variables to values. Lookup on a Var bound to another
// Var is recursive.
typealias Substitution = [Var: Term]
extension Substitution {
// Recursively looks up `u` in self until we either reach a primitive value or
// an unbound variable.
func walk(_ u: Term) -> Term {
if case .var(let u) = u, let v = self[u] {
return walk(v)
} else {
return u
}
}
func extend(_ x: Var, _ v: Term) -> Substitution {
return merging([x: v]) { current, new in new }
}
}
// A set of bindings and the index of the next variable that will be allocated.
//
// One valid result of evaluating a logic expression. In other words, one answer to
// the question you're asking.
struct State: CustomStringConvertible {
let sub: Substitution
var nfresh: Int
static let empty = State(sub: Substitution(), nfresh: 0)
var description: String {
return "{\(sub) \(nfresh)}"
}
mutating func makeVar() -> Term {
defer { nfresh += 1 }
return .var(Var(nfresh))
}
}
// A (potentially infinite) linked list of States. This is the result of evaluating a
// miniKanren program. Infinite streams are represented lazily. A non-lazy (mature)
// stream is just a cons cell. A lazy (immature) stream is a thunk – a function that
// returns a stream.
enum Stream {
case empty
indirect case mature(State, Stream)
case immature(() -> Stream)
// A stream with a single element.
static func unit(_ state: State) -> Stream {
return .mature(state, .empty)
}
// Concatenates two streams. Non-lazy streams are concatenated as you'd expect.
// Upon encountering a lazy stream, `other` is inserted before the rest of the
// (potentially infinite) lazy stream. IIRC this interleaves the results so that
// solutions are returned breath-first instead of depth-first.
func append(_ other: Stream) -> Stream {
switch self {
case .empty:
return other
case .immature(let f):
return .immature { other.append(f()) }
case .mature(let state, let rest):
return .mature(state, rest.append(other))
}
}
// Similar to flat-map or the monadic bind operator. Evaluates a Goal (a logic
// expression) with each of the intermediate States in self, returning a new
// stream with the evaluations results. The resulting stream can be shorter
// (if `g` represents an "and" expression) or longer (if `g` is an "or").
func appendMap(_ g: @escaping Goal) -> Stream {
switch self {
case .empty:
return .empty
case .immature(let f):
return .immature { f().appendMap(g) }
case .mature(let state, let rest):
return g(state).append(rest.appendMap(g))
}
}
}
extension Stream: Sequence {
struct Iterator: IteratorProtocol {
var stream: Stream
mutating func next() -> Stream? {
switch stream {
case .mature(_, let next):
defer { stream = next }
return stream
case .immature(_):
defer { stream = .empty }
return stream
case .empty:
return nil
}
}
}
public func makeIterator() -> Iterator {
return Iterator(stream: self)
}
}
extension Stream: CustomStringConvertible {
var description: String {
let elements = compactMap { stream in
switch stream {
case .mature(let state, _):
return String(describing: state)
case .immature(_):
return "..."
case .empty:
return nil
}
}.joined(separator: " ")
return "(\(elements))"
}
}
// Attempts to bind `u` and `v` to the same underlying value in the context of
// an existing set of bindings. Logic variables are recursively evaluated until
// they either evaluate to a concrete value or are found to be not yet bound.
//
// This operation is the basis of type inference in many programming languages.
//
// Upon success, returns a new set of bindings where `u` and `v` will evaluate
// to the same value. Returns nil on failure.
//
// Succeeds in any of the following conditions:
// 1. `u` and `v` evaluate to the same unbound variable.
// 2. Either `u` or `v` evaluates to an unbound variable. In this case, returns
// a new set of bindings where `u` and `v` are equivalent.
// 3. `u` and `v` both evaluate to lists that can be unified pairwise.
// 4. `u` and `v` evaluate to the the same primitive value.
func unify(_ u: Term, _ v: Term, _ s: Substitution) -> Substitution? {
let u = s.walk(u)
let v = s.walk(v)
if case .var(let u) = u, case .var(let v) = v, u == v {
return s
} else if case .var(let u) = u {
return s.extend(u, v)
} else if case .var(let v) = v {
return s.extend(v, u)
} else if case .pair(let ua, let ud) = u,
case .pair(let va, let vd) = v {
guard let s = unify(ua, va, s) else {
return nil
}
return unify(ud, vd, s)
} else if u == v {
return s
} else {
return nil
}
}
// A logic expression. Represented as a function from a set of bindings to a
// (possibly infinite) stream of bindings that satisfy the logic expression.
//
// Any function that returns a Goal is an operator. A tool that lets you build
// up larger and larger logic expressions.
typealias Goal = (State) -> Stream
// The most primitive logic expression. Can be read as "matches." If `u` and `v`
// are equal, return a single result, otherwise return no results.
infix operator =~: ComparisonPrecedence
func =~ (_ u: Term, _ v: Term) -> Goal {
return { state in
if let sub = unify(u, v, state.sub) {
return Stream.unit(State(sub: sub, nfresh: state.nfresh))
} else {
return .empty
}
}
}
// Introduce a new ("fresh") variable, which starts off unbound.
func callFresh(f: @escaping (Term) -> Goal) -> Goal {
return { state in
var state = state
let g = f(state.makeVar())
return g(state)
}
}
// Conjunction. The primitive "and". Upon evaluation, this operator matches
// (i.e. "succeeds") if both of its operands match.
func conj_(_ g1: @escaping Goal, _ g2: @escaping Goal) -> Goal {
return { state in
g1(state).appendMap(g2)
}
}
// Disjunction. The primitive "or". Matches if either operands match.
func disj_(_ g1: @escaping Goal, _ g2: @escaping Goal) -> Goal {
return { state in
g1(state).append(g2(state))
}
}
// miniKanren – a higher level relational language. Has the same semantics as µKanren,
// but nicer to use.
// Sleep. An operator that evaluates its operand lazily.
func zzz(_ g: @escaping @autoclosure () -> Goal) -> Goal {
return { state in
return .immature {
g()(state)
}
}
}
// A lazy "and" operator.
func && (_ g1: @escaping @autoclosure () -> Goal, _ g2: @escaping @autoclosure () -> Goal) -> Goal {
conj_(zzz(g1()), zzz(g2()))
}
// A lazy "or" operator.
func || (_ g1: @escaping @autoclosure () -> Goal, _ g2: @escaping @autoclosure () -> Goal) -> Goal {
disj_(zzz(g1()), zzz(g2()))
}
// Primitive expressions that always succeed or fail, regardless of the bindings they're
// evaluated in. Think `true` and `false`.
let succeed: Goal = { state in .unit(state) }
let fail: Goal = { state in .empty }
// Variadic "and".
func conj(@ConjBuilder f: () -> Goal) -> Goal {
f()
}
@resultBuilder struct ConjBuilder {
static func buildBlock(_ goals: Goal...) -> Goal {
goals.reduce(succeed) { x, y in x && y }
}
}
// Variadic "or".
func disj(@DisjBuilder f: () -> Goal) -> Goal {
f()
}
@resultBuilder struct DisjBuilder {
static func buildBlock(_ goals: Goal...) -> Goal {
goals.reduce(fail) { x, y in x || y }
}
}
// If-then-else. Works just like you expect.
func ifte(_ g1: @escaping Goal, _ g2: @escaping Goal, _ g3: @escaping Goal) -> Goal {
return { state in
var stream = g1(state)
while true {
switch stream {
case .empty:
return g3(state)
case .mature(_, _):
return stream.appendMap(g2)
case .immature(let f):
stream = f()
}
}
}
}
// if [else if]* else?. Easier to write than a bunch of nested if-then-elses.
// Short circuits, just like imperative if-then-elses. Only the first matching
// branch is evaluated and contributes values.
func conda(@CondaBuilder f: () -> Goal) -> Goal {
f()
}
@resultBuilder struct CondaBuilder {
static func buildBlock(_ goals: [Goal]...) -> Goal {
makeGoal(goals)
}
static func makeGoal(_ goals: [[Goal]]) -> Goal {
var g = succeed
for line in goals.reversed() {
if line.count > 0 {
g = ifte(
line[0],
line.dropFirst().reduce(succeed) { a, b in a && b },
g)
}
}
return g
}
}
// Matches the first result of `g`, discarding the rest.
func once(_ g: @escaping Goal) -> Goal {
return { state in
var stream = g(state)
while true {
switch stream {
case .empty:
return .empty
case .mature(let state, _):
return .unit(state)
case .immature(let f):
stream = f()
}
}
}
}
// Like conda, but the first matching predicate only contributes a single value,
// discarding the rest. E.g if the matching predicate of a conda is `x =~ 3 || x =~ 4`,
// two values for x (3 and 4) will be contributed to the result. If the same predicate
// is used with condu, only 3 will be contributed to the result.
func condu(@ConduBuilder f: () -> Goal) -> Goal {
f()
}
@resultBuilder struct ConduBuilder {
static func buildBlock(_ goals: [Goal]...) -> Goal {
let goals: [[Goal]] = goals.map { line in
if line.count > 0 {
return [once(line[0])] + line.dropFirst()
} else {
return line
}
}
return CondaBuilder.makeGoal(goals)
}
}
extension Stream {
// Eagerly evaluate the entire stream into an array. Calling takeAll on an
// infinite stream will never return.
func takeAll() -> [State] {
switch self {
case .empty:
return []
case .mature(let state, let rest):
return [state] + rest.takeAll()
case .immature(let f):
return f().takeAll()
}
}
// Evaluate the first n elements in the stream. Guaranteed to return even
// if self is infinite.
func take(_ n: Int) -> [State] {
if n == 0 {
return []
}
switch self {
case .empty:
return []
case .mature(let state, let rest):
return [state] + rest.take(n-1)
case .immature(let f):
return f().take(n)
}
}
}
// Creates a sym – a printable representation of an unbound variable.
func reifyName(_ n: Int) -> Term {
return .sym("_\(n)")
}
extension Substitution {
// Like walk, looks up `v` in self, but also walks lists so that the only
// variables contained in walkAll's return value are unbound.
func walkAll(_ v: Term) -> Term {
let v = walk(v)
if case .var(_) = v {
return v
} else if case .pair(let value, let rest) = v {
return .pair(walkAll(value), walkAll(rest))
} else {
return v
}
}
// Recursively looks up `v`, returning a new set of bindings where every
// unbound variable in self is bound to a sym, a printable representation
// of itself (e.g. _0, _1, _2, etc.)
func reifyS(_ v: Term) -> Substitution {
let v = walk(v)
if case .var(let v) = v {
let n = count
let name = reifyName(n)
return extend(v, name)
} else if case .pair(let value, let rest) = v {
let r = reifyS(value)
return r.reifyS(rest)
} else {
return self
}
}
// Recursively looks up `v`, replacing all the resulting unbound vars with
// their printable representations.
func reify(_ v: Term) -> Term {
let v = walkAll(v)
let names = Substitution().reifyS(v)
return names.walkAll(v)
}
}
// Introduces an unbound variable, evaluated in the context of multiple
// sub-expressions, which are all anded together.
func fresh(@ConjBuilder f: @escaping (Term) -> Goal) -> Goal {
callFresh { a in f(a) }
}
// Forces the evaluation of a series of sub-expressions, in the context of a single
// unbound variable. All sub-expressions are anded together. Stops evaluation after
// the first n results. If n is nil, returns all results. If n is nil and there are
// an infinite number of solutions, never returns.
func run(n: Int? = nil, @ConjBuilder f: (Term) -> Goal) -> [Term] {
var state = State.empty
let a = state.makeVar()
if let n = n {
return f(a)(state).take(n).map { $0.sub.reify(a) }
} else {
return f(a)(state).takeAll().map { $0.sub.reify(a) }
}
}
// For convenience
// Introduce multiple unbound variables at once.
func fresh(@ConjBuilder f: @escaping (Term, Term) -> Goal) -> Goal {
fresh { a in fresh { b in f(a, b) }}
}
func fresh(@ConjBuilder f: @escaping (Term, Term, Term) -> Goal) -> Goal {
fresh { a, b in fresh { c in f(a, b, c) }}
}
func fresh(@ConjBuilder f: @escaping (Term, Term, Term, Term) -> Goal) -> Goal {
fresh { a, b, c in fresh { d in f(a, b, c, d) }}
}
func fresh(@ConjBuilder f: @escaping (Term, Term, Term, Term, Term) -> Goal) -> Goal {
fresh { a, b, c, d in fresh { e in f(a, b, c, d, e) }}
}
func fresh(@ConjBuilder z: @escaping (Term, Term, Term, Term, Term, Term) -> Goal) -> Goal {
fresh { a, b, c, d, e in fresh { f in z(a, b, c, d, e, f) }}
}
func fresh(@ConjBuilder z: @escaping (Term, Term, Term, Term, Term, Term, Term) -> Goal) -> Goal {
fresh { a, b, c, d, e, f in fresh { g in z(a, b, c, d, e, f, g) }}
}
func fresh(@ConjBuilder z: @escaping (Term, Term, Term, Term, Term, Term, Term, Term) -> Goal) -> Goal {
fresh { a, b, c, d, e, f, g in fresh { h in z(a, b, c, d, e, f, g, h) }}
}
func fresh(@ConjBuilder z: @escaping (Term, Term, Term, Term, Term, Term, Term, Term, Term) -> Goal) -> Goal {
fresh { a, b, c, d, e, f, g, h in fresh { i in z(a, b, c, d, e, f, g, h, i) }}
}
func fresh(@ConjBuilder z: @escaping (Term, Term, Term, Term, Term, Term, Term, Term, Term, Term) -> Goal) -> Goal {
fresh { a, b, c, d, e, f, g, h, i in fresh { j in z(a, b, c, d, e, f, g, h, i, j) }}
}
func fresh(@ConjBuilder z: @escaping (Term, Term, Term, Term, Term, Term, Term, Term, Term, Term, Term) -> Goal) -> Goal {
fresh { a, b, c, d, e, f, g, h, i, j in fresh { k in z(a, b, c, d, e, f, g, h, i, j, k) }}
}
// Force evaluation, introducing multiple unbound variables at once.
func run(n: Int? = nil, @ConjBuilder f: (Term, Term) -> Goal) -> [Term] {
var state = State.empty
let a = state.makeVar()
let b = state.makeVar()
if let n = n {
return f(a, b)(state).take(n).map { $0.sub.reify([a, b]) }
} else {
return f(a, b)(state).takeAll().map { $0.sub.reify([a, b]) }
}
}
func run(n: Int? = nil, @ConjBuilder f: (Term, Term, Term) -> Goal) -> [Term] {
var state = State.empty
let a = state.makeVar()
let b = state.makeVar()
let c = state.makeVar()
if let n = n {
return f(a, b, c)(state).take(n).map { $0.sub.reify([a, b, c]) }
} else {
return f(a, b, c)(state).takeAll().map { $0.sub.reify([a, b, c]) }
}
}
func run(n: Int? = nil, @ConjBuilder f: (Term, Term, Term, Term) -> Goal) -> [Term] {
var state = State.empty
let a = state.makeVar()
let b = state.makeVar()
let c = state.makeVar()
let d = state.makeVar()
if let n = n {
return f(a, b, c, d)(state).take(n).map { $0.sub.reify([a, b, c, d]) }
} else {
return f(a, b, c, d)(state).takeAll().map { $0.sub.reify([a, b, c, d]) }
}
}
func run(n: Int? = nil, @ConjBuilder f: (Term, Term, Term, Term, Term) -> Goal) -> [Term] {
var state = State.empty
let a = state.makeVar()
let b = state.makeVar()
let c = state.makeVar()
let d = state.makeVar()
let e = state.makeVar()
if let n = n {
return f(a, b, c, d, e)(state).take(n).map { $0.sub.reify([a, b, c, d, e]) }
} else {
return f(a, b, c, d, e)(state).takeAll().map { $0.sub.reify([a, b, c, d, e]) }
}
}
// Now we're in user space. From here on out, there are no more new primitives. Now we're just
// building larger logic expressions out of smaller ones. Let's play!
// Start by defining some useful relations. Relations are functions that end in "o" and return
// a Goal. You can think of a relation as a named sub-expression. Nested calls to relations
// build up a tree of logic expressions to be evaluated.
// Each relation is the equivalent of some normal function in the imperative world. Relations
// corresponding to imperative predicates – functions that return a Bool, like isNull(_:) –
// take the same number of arguments. Relations corresponding to functions with non-Bool
// return values take an extra argument representing their return value. E.g. append(left, right)
// becomes appendo(left, right, result).
// isNull
func nullo(_ x: Term) -> Goal {
x =~ []
}
// pair = cons(a, b)
func conso(_ a: Term, _ d: Term, _ p: Term) -> Goal {
p =~ .pair(a, d)
}
// first = car(list)
func caro(_ l: Term, _ a: Term) -> Goal {
fresh { d in
conso(a, d, l)
}
}
// rest = cdr(list)
func cdro(_ l: Term, _ d: Term) -> Goal {
fresh { a in
conso(a, d, l)
}
}
// l.isList && list.count == 1
func singletono(_ l: Term) -> Goal {
fresh { d in
cdro(l, d)
nullo(d)
}
}
// result = append(list1, list2)
func appendo(_ l: Term, _ t: Term, _ out: Term) -> Goal {
disj {
nullo(l) && t =~ out
fresh { a, d, res in
conso(a, d, l)
conso(a, res, out)
appendo(d, t, res)
}
}
}
// With that out of the way, let's ask some questions.
// "Append the lists (1 "foo") and ("bar" 4)"
var r = run { a in
appendo([1, "foo"], ["bar", 4], a)
}
print(r) // => [(1 "foo" "bar" 4)]
// "Append the numbers 1 and 2"
//
// Returns no results because appendo is only defined for lists.
r = run { a in
appendo(1, 2, a)
}
print(r) // => []
// "Give me all pairs of lists that append to (1 "foo" "bar" 4)"
r = run { a, b in
appendo(a, b, [1, "foo", "bar", 4])
}
print(r) // => [(() (1 "foo" "bar" 4)), ((1) ("foo" "bar" 4)), ((1 "foo") ("bar" 4)), ((1 "foo" "bar") (4)), ((1 "foo" "bar" 4) ())]
// "Give me all possible evaluations of `appendo`"
//
// We must specify `n` because there are infinite answers to this query.
r = run(n: 5) { q, r, s in
appendo(q, r, s)
}
print(r) // => [(() _0 _0), ((_0) _1 (_0 . _1)), ((_0 _1) _2 (_0 _1 . _2)), ((_0 _1 _2) _3 (_0 _1 _2 . _3)), ((_0 _1 _2 _3) _4 (_0 _1 _2 _3 . _4))]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment