miniKanren

miniKanren.swift

// µKanren

struct Var: Equatable, Hashable {
    let i: Int
    
    init(_ i: Int) {
        self.i = i
    }
}

enum Term: Equatable {
    case none
    case int(Int)
    case double(Double)
    case string(String)
    case bool(Bool)
    indirect case pair(Term, Term)
    case `var`(Var)
    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
            }
            
            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)
    }
}

typealias Substitution = [Var: Term]

extension Substitution {
    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 }
    }
}

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))
    }
}

enum Stream {
    case empty
    indirect case mature(State, Stream)
    case immature(() -> Stream)

    static func unit(_ state: State) -> Stream {
        return .mature(state, .empty)
    }
    
    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))
        }
    }
    
    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))"
    }
}

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
    }
}

typealias Goal = (State) -> Stream

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
        }
    }
}

func callFresh(f: @escaping (Term) -> Goal) -> Goal {
    return { state in
        var state = state
        let g = f(state.makeVar())
        
        return g(state)
    }
}

func conj_(_ g1: @escaping Goal, _ g2: @escaping Goal) -> Goal {
    return { state in
        g1(state).appendMap(g2)
    }
}

func disj_(_ g1: @escaping Goal, _ g2: @escaping Goal) -> Goal {
    return { state in
        g1(state).append(g2(state))
    }
}



// miniKanren

func zzz(_ g: @escaping @autoclosure () -> Goal) -> Goal {
    return { state in
        return .immature {
            g()(state)
        }
    }
}

func && (_ g1: @escaping @autoclosure () -> Goal, _ g2: @escaping @autoclosure () -> Goal) -> Goal {
    conj_(zzz(g1()), zzz(g2()))
}

func || (_ g1: @escaping @autoclosure () -> Goal, _ g2: @escaping @autoclosure () -> Goal) -> Goal {
    disj_(zzz(g1()), zzz(g2()))
}

let succeed: Goal = { state in .unit(state) }
let fail: Goal = { state in .empty }

@_functionBuilder struct ConjBuilder {
    static func buildBlock(_ goals: Goal...) -> Goal {
        goals.reduce(succeed) { x, y in x && y }
    }
}

@_functionBuilder struct DisjBuilder {
    static func buildBlock(_ goals: Goal...) -> Goal {
        goals.reduce(fail) { x, y in x || y }
    }
}

func conj(@ConjBuilder f: () -> Goal) -> Goal {
    f()
}

func disj(@DisjBuilder f: () -> Goal) -> Goal {
    f()
}

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()
            }
        }
    }
}


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()
            }
        }
    }
}


@_functionBuilder 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
    }
}

func conda(@CondaBuilder f: () -> Goal) -> Goal {
    f()
}

@_functionBuilder 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)
    }
}

func condu(@ConduBuilder f: () -> Goal) -> Goal {
    f()
}

extension Stream {
    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()
        }
    }
    
    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)
        }
    }
}

func reifyName(_ n: Int) -> Term {
    return .sym("_\(n)")
}

extension Substitution {
    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
        }
    }
    
    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
        }
    }
    
    func reify(_ v: Term) -> Term {
        let v = walkAll(v)
        let names = Substitution().reifyS(v)

        return names.walkAll(v)
    }
}

func fresh(@ConjBuilder f: @escaping (Term) -> Goal) -> Goal {
    callFresh { a in f(a) }
}


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

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) }}
}



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


// Actual playground

func nullo(_ x: Term) -> Goal {
    x =~ []
}

func conso(_ a: Term, _ d: Term, _ p: Term) -> Goal {
    p =~ .pair(a, d)
}

func caro(_ l: Term, _ a: Term) -> Goal {
    fresh { d in
        conso(a, d, l)
    }
}

func cdro(_ l: Term, _ d: Term) -> Goal {
    fresh { a in
        conso(a, d, l)
    }
}

func singletono(_ l: Term) -> Goal {
    fresh { d in
        cdro(l, d)
        nullo(d)
    }
}

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)
        }
    }
}


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