miniKanren

miniKanren.swift

// 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))]