sharing.swift

protocol Expr {
    static func constant(_: Int) -> Self
    static func variable(_: String) -> Self
    static func add(_: Self, _: Self) -> Self
}

indirect enum Exp {
    case _const(Int)
    case _variable(String)
    case _add(Exp, Exp)
}

extension Exp: Expr {
    static func constant(_ x: Int) -> Exp {
        return _const(x)
    }
    
    static func variable(_ x: String) -> Exp {
        return _variable(x)
    }
    
    static func add(_ l: Exp, _ r: Exp) -> Exp {
        return _add(l, r)
    }
}

func exp_b<E: Expr>() -> E {
    let exp_a: E = .add(.constant(10), .variable("i1"))
    return .add(exp_a, .variable("i2"))
}

func mul<E: Expr>(_ x: Int, _ y: E) -> E {
    switch x {
    case 0: return .constant(0)
    case 1: return y
    case let n where n % 2 == 0:
        return mul(n/2, .add(y,y))
    case let n:
        return E.add(y, mul(n-1, y))
    }
}

func exp_mul4<E: Expr>() -> E {
    return mul(4, .variable("i1"))
}

dump(exp_b() as Exp)

typealias REnv = [String:Int]
struct R { let unR: (REnv) -> Int }
extension R: Expr {
    static func constant(_ x: Int) -> R {
        return R { _ in x }
    }
    static func variable(_ n: String) -> R {
        return R { env in env[n]! }
    }
    static func add(_ l: R, _ r: R) -> R {
        return R { l.unR($0) + r.unR($0) }
    }
}

(exp_mul4() as R).unR(["i1":5])

typealias NodeID = Int
enum Node {
    case nconst(Int)
    case nvar(String)
    case nadd(NodeID,NodeID)
}
extension Node: Equatable {
    static func ==(lhs: Node, rhs: Node) -> Bool {
        switch (lhs, rhs) {
        case (.nconst(let l), .nconst(let r)): return l == r
        case (.nvar(let l), .nvar(let r)): return l == r
        case (.nadd(let l), .nadd(let r)): return l == r
        default: return false
        }
    }
}

struct Bimap<A> where A: Equatable {
    private var storage: [A] = []
    
    subscript(val value: A) -> Int? {
        return storage.index(where: { $0 == value})
    }
    subscript(key key: Int) -> A {
        return storage[key]
    }
    
    mutating func insert(_ value: A) -> Int {
        storage.append(value)
        return storage.count - 1
    }
}

struct State<S,R> {
    let modify: (inout S) -> R
    func run(_ s: S) -> (R, S) {
        var x = s
        let result = modify(&x)
        return (result, x)
    }
    init(modify: @escaping (inout S) -> R) {
        self.modify = modify
    }
}

typealias DAG = Bimap<Node>

struct N {
    let unN: State<DAG, NodeID>
}

extension N {
    var run: (NodeID, DAG) {
        return unN.run(DAG())
    }
}

extension Node {
    var hashcons: State<DAG,NodeID> {
        return State { (s: inout DAG) in
            s[val: self] ?? s.insert(self)
        }
    }
}

extension N: Expr {
    static func constant(_ x: Int) -> N {
        return N(unN: Node.nconst(x).hashcons)
    }
    static func variable(_ s: String) -> N {
        return N(unN: Node.nvar(s).hashcons)
    }
    static func add(_ e1: N, _ e2: N) -> N {
        return N(unN: State { (state: inout DAG) in
            let h1 = e1.unN.modify(&state)
            let h2 = e2.unN.modify(&state)
            return Node.nadd(h1, h2).hashcons.modify(&state)
        })
    }
}

dump((exp_mul4() as N).run)