Playing with Swift Pattern Matching for Mathematical Expressions (run in playground)

SwiftExpr.swift

import Cocoa

//: # An exercise in specifying and simplifying Maths functions using Swift
/*: This is cut-down version of some code I've been working on to parse, simplify and display mathematical expressions. The source below shows some of the power of pattern matching and (recursive) enumerated types within Swift in defining and evaluating the expression domain. Note, it doesn't cover all cases but is intended to give a flavour of what can be achieved.

    No optimisations applied
*/

//: # Supporting functions
//: ### Function Pipeline
infix operator |> { associativity left precedence 140 }
public func |> <T,U> (@autoclosure left: () -> T,right: T -> U) -> U {
    return right(left())
}

//: ### Add/Update element in a dictionary returning an updated copy
extension Dictionary {
    func add(k: Key, v : Value? -> Value) -> Dictionary<Key,Value> {
        var copy = self
        copy[k] = v(self[k])
        return copy
    }
}

//: # Expression Domain
//: A simple (incomplete) set of options representing mathmatical operations.
/*: 
- Integer numbers  **Num(10)**
- Named variables  **Var("x")**
- Negative expressions **Neg(Num(1))**
- Additions **Add([Num(10),Num(20),Var("x")])**
- Subtractions **Sub(Num(10),[Num(20)])**  
    You'd probably expect subtraction to be Sub(Expr,Expr) but there is a good reason I split into Sub(Expr,[Expr]) - consider the expressions:


    let x = 10 - 20 - 30
    let y = 10 - (20 - 30)
    
    - first result is -40
    - second result is 20


- Multiplications **Prod(Num(10),Num(10)]**
- Divisions **Frac(Num(10),Num(2))**
- Powers **Pow(Var("x"),Num(2))**
*/
indirect enum Expr   {
    case Num(Int)
    case Var(String)
    case Neg(Expr)
    case Add([Expr])
    case Sub(Expr, [Expr])
    case Prod(Expr,Expr)
    case Frac(Expr,Expr)
    case Pow(Expr,Int)
    case NaN
}

/*: #### provided to help display within the playground (useful display of string interpolation)
Does have scope for improvement around precedence but sufficient for test purposes
*/
extension Expr : CustomDebugStringConvertible  {
    var debugDescription : String {
        func exprStr(e: Expr) -> String {
            switch e {
            case .Var(let vn)           : return vn
            case .Num(let n)            : return "\(n)"
            case .Neg(let e1)           : return "-" + exprStr(e1)
            case .Add(let exprs)        : return exprs.map({exprStr($0)}).joinWithSeparator("+")
            case .Sub(let s1, let sl)   : return exprStr(s1) + "-" + sl.map({exprStr($0)}).joinWithSeparator("-")
            case .Prod(let e1, let e2)  :
                switch (e1,e2) {
                case (.Var(let v), .Num(let n)) : return "(\(n)\(v))"
                case (.Num(let n), .Var(let v)) : return "(\(n)\(v))"
                case (.Var(let v), let e) : return "\(v)\(exprStr(e))"
                case (let e, .Var(let v)) : return "\(v)\(exprStr(e))"
                case _ : return "(\(exprStr(e1)) * \(exprStr(e2)))"
                }
            case .Frac(let e1, let e2)  : return "(\(exprStr(e1)) / \(exprStr(e2)))"
            case .Pow(let e1, let p)    : return "\(exprStr(e1))^\(p)"
            case .NaN                   : return "NaN"
            }
        }
        
        return exprStr(self)
    }
}

//: # Simplification Functions

func negate(e: Expr) -> Expr {
    switch e {
    case .Num(let n)            : return .Num(-n)
    case .Var(_)                : return .Neg(e)
    case .Neg(let ex)           : return ex
    case .Add(let exprs)        : return .Add(exprs.map(negate))
    case .Prod(let e1, let e2)  : return .Prod(negate(e1), e2)
    case .Frac(let e1, let e2)  : return .Frac(negate(e1), e2)
    case let ex                 : return .Neg(ex)
    }
}

func addExprs(e: Expr) -> [Expr] {
    guard case .Add(let exprs) = e else {
        return [e]
    }
    return exprs
}

/*: ## Expression Simplify

An implementation of a function to take an expression and simplify down to (reasonably) simple form.

*/
func simplify(e: Expr) -> Expr {
    switch e {
    case .Num(let n)    : return .Num(n)
    case .Var(let v)    : return .Var(v)
    case .Neg(let ex)   : return negate(simplify(ex))
    case .Add(let exprs):
        /*:
            - take all the 'add' expressions and simplify and then flatten to a single [Expr]
            - fold through the resulting list 
                - adding up all the numbers
                - creating a 'variable' dictionary (with counts)
                - capturing all 'other' exprs we can't add up or aren't vars
            - just numbers => result is the total
            - no numbers but other expr => add(products + expr)
            - otherwise => add(total, products, expr)
        */
        
        let simplified = exprs.map(simplify).map(addExprs).flatMap {$0}
        let (total,vars,el) = simplified.reduce((0,[:],[]), combine: { (acc , element) -> (Int,[String:Int],[Expr]) in
            switch element {
            case .Num(let n)                        : return (acc.0 + n, acc.1, acc.2)
            case .Var(let vn)                       : return (acc.0, acc.1.add(vn) {($0 ?? 0) + 1}, acc.2)
            case .Prod(.Num(let n), .Var(let vn))   : return (acc.0, acc.1.add(vn) {($0 ?? 0) + n}, acc.2)
            case .Prod(.Var(let vn),.Num(let n))    : return (acc.0, acc.1.add(vn) {($0 ?? 0) + n}, acc.2)
            case _ : return (acc.0, acc.1, acc.2 + [element])
            }
        })
        
        let products = vars.map {(k,v) in v > 1 ? Expr.Prod(Expr.Num(v),Expr.Var(k)) : Expr.Var(k)}
        switch (total,products.count,el.count) {
        case (let t, 0, 0) where t > 0                  : return .Num(t)
        case (0, let p, let x)  where p > 0 || x > 0    : return .Add(products + el)
        case _                                          : return .Add([.Num(total)] + products + el)
        }
        
    case .Sub(let e1, let exprs): return simplify(.Add([e1] + exprs.map(Expr.Neg)))
    case .Prod(let e1, let e2):
        let s1 = simplify(e1)
        let s2 = simplify(e2)
        switch (s1,s2) {
        case (.Num(0), _)                               : return .Num(0)
        case (_, .Num(0))                               : return .Num(0)
        case (.Num(1), let ex)                          : return ex
        case (let ex, .Num(1))                          : return ex
        case (.Num(let n1) , .Num(let n2))              : return .Num(n1 * n2)
        case (.Num(let n1) , .Prod(.Num(let n), let e)) : return .Add([.Num(n1 * n), .Prod(.Num(n1), e)])
        case (.Var(let v1), .Var(let v2))
            where v1 == v2                              : return .Pow(.Var(v1), 2)
        case (.Var(let v1), .Pow(.Var(let v2), let p))
            where v1 == v2                              : return .Pow(.Var(v1), p+1)
        case (let e1, let e2)                           : return .Prod(e1,e2)
        }
    case .Frac(let e1, let e2):
        let s1 = simplify(e1)
        let s2 = simplify(e2)
        switch (s1,s2) {
        case (.Num(0), _)                                   : return .Num(0)
        case (let ex, .Num(1))                              : return ex
        case (_ , .Num(0))                                  : return .NaN
        case (.Num(let n1), .Num(let n2))                   : return .Num(n1 / n2)
        case (.Num(let n1), .Frac(.Num(let n2), let ex))    : return .Prod(.Frac(.Num(n1), .Num(n2)), ex)
        case (.Num(let n1), .Frac(let ex, .Num(let n2)))    : return .Frac(.Prod(.Num(n1), .Num(n2)), e)
        case (let e1, let e2)                               : return .Frac(e1,e2)
        }
    case .Pow(let ex, 1)        : return simplify(ex)
    case .Pow(let ex, let n)    :
        switch simplify(ex) {
        case .Num(let v) : return .Num(Int(pow(CGFloat(v),CGFloat(n))))
        case let sp : return .Pow(sp, n)
        }
    case .NaN                   : return e
    }
}

let simple : Expr -> Expr = {$0 |> simplify |> simplify  } // multi-pass simplify

//: #### Sample Test Expressions
let a1 = Expr.Add([Expr.Num(10), Expr.Prod(Expr.Num(2), Expr.Var("x")), Expr.Num(20), Expr.Num(30), Expr.Var("x"), Expr.Add([Expr.Num(1), Expr.Var("y")])])
let a2 = Expr.Prod(.Num(3), .Var("x"))
let a3 = Expr.Add([a1,a2])
let f1 = Expr.Frac(a1, Expr.Pow(a2, 2))
let r1 = simple(a3)
let r2 = simple(f1)

let xx = Expr.Add([Expr.Prod(Expr.Var("x"), Expr.Var("y")), Expr.Var("z")])
let rx = simple(xx)


// subtraction
let s1 = Expr.Sub(Expr.Var("x"), [Expr.Num(20),Expr.Num(5),Expr.Var("x")])
let rs1 = simple(s1)
let x = 10 - 20 - 30    // why sub is modelled as an [Expr]
let y = 10 - (20 - 30)

// division
let d1 = Expr.Frac(f1, Expr.Num(2))
let rd1nc = simple(d1)

let d1zero = Expr.Frac(f1, Expr.Num(0))
let rd1nczero = simple(d1zero)
let d2zero = Expr.Add([Expr.Num(10), Expr.Frac(Expr.Num(2), Expr.Num(0)), Expr.Num(20), Expr.Num(30), Expr.Var("x"), Expr.Add([Expr.Num(1), Expr.Var("y")])])
let rd2nczero = simple(d2zero) // should fix the issue where end result should be NaN here.


// Powers
let p1 = Expr.Pow(Expr.Num(10), 2)
let p2 = Expr.Add([Expr.Var("x"),p1,Expr.Var("x")])
let rp1 = simple(p2)

// Products
let pr1 = Expr.Prod(Expr.Num(2), Expr.Var("x"))
let pr2 = Expr.Prod(Expr.Num(2), pr1)
let rpr1 = simple(pr2)

let pr3 = Expr.Prod(Expr.Var("x"), Expr.Var("y"))
let pr4 = Expr.Prod(Expr.Num(2), pr3)
let rpr2 = simple(pr4)


/*: # And just for fun
 note this is easily breakable (forced optionals)
*/
func exprType<Z>(v : Z) -> Expr? {
    switch v {
    case let x as IntegerLiteralType : return Expr.Num(x)
    case let x as StringLiteralType : return Expr.Var(x)
    case let x as Expr : return x
    case _ : return nil
    }
}

infix operator |+| { associativity left precedence 140 }
func |+| <T,U> (@autoclosure l: () -> T, @autoclosure  r: () -> U) -> Expr {
    return Expr.Add([exprType(l()),exprType(r())].flatMap({$0}))
}

infix operator |-| { associativity left precedence 140 }
func |-| <T,U> (@autoclosure l: () -> T, @autoclosure  r: () -> U) -> Expr {
    return Expr.Sub(exprType(l())!, [exprType(r())!]) // only allow a single sub
}

infix operator |*| { associativity left precedence 160 }
func |*| <T,U> (@autoclosure l: () -> T, @autoclosure  r: () -> U) -> Expr {
    return Expr.Prod(exprType(l())!,exprType(r())!)
}

infix operator |/| { associativity left precedence 160 }
func |/| <T,U> (@autoclosure l: () -> T, @autoclosure  r: () -> U) -> Expr {
    return Expr.Frac(exprType(l())!,exprType(r())!)
}

prefix operator !! {}
prefix func !! (@autoclosure  r: () -> Int) -> Expr {
    return Expr.Num(r())
}

prefix operator ^= {}
prefix func ^= (@autoclosure  r: () -> Expr) -> Expr {
    return simple(r())
}

let result1 = ^=(100 |*| 2 |+| "x" |+| "x" |-| 10 |+| (2 |*| "x"))