HOAS in finally tagless style in Swift 3

finally-tagless-hoas.swift

// A (sadly untyped) adaptation of http://okmij.org/ftp/tagless-final/CB98.hs

protocol EDSL : ExpressibleByIntegerLiteral {
  static func add(_ lhs: Self, _ rhs: Self) -> Self?
  static func mul(_ lhs: Self, _ rhs: Self) -> Self?
  static func lam(_ body: @escaping (Self) -> Self) -> Self
  static func app(_ fn: Self, _ arg: Self) -> Self?
}

extension EDSL {
  static func let_(_ x: Self, _ y: @escaping (Self) -> Self?) -> Self {
    return app(lam{z in y(z)!}, x)!
  }
}

func +<Exp : EDSL>(lhs : Exp, rhs: Exp) -> Exp {
  return Exp.add(lhs, rhs)!
}
func *<Exp : EDSL>(lhs : Exp, rhs: Exp) -> Exp {
  return Exp.mul(lhs, rhs)!
}


enum SValue {
  case Int(Int);
  case Lam((@escaping () -> SValue) -> SValue);
}

extension SValue : EDSL {
  init(integerLiteral : Int) {
    self = .Int(integerLiteral)
  }
  static func add(_ lhs: SValue, _ rhs: SValue) -> SValue? {
    switch (lhs, rhs) {
    case (.Int(let x), .Int(let y)):
      print("Adding")
      return Int(x + y)
    default:
      return nil
    }
  }
  static func mul(_ lhs: SValue, _ rhs: SValue) -> SValue? {
    switch (lhs, rhs) {
    case (.Int(let x), .Int(let y)):
      print("Multiplying")
      return Int(x * y)
    default:
      return nil
    }
  }
  static func lam(_ body: @escaping (SValue) -> SValue) -> SValue {
    return Lam { v in body(v()) }
  }
  static func app(_ fn: SValue, _ arg: SValue) -> SValue? {
    switch fn {
    case .Lam(let f):
      return f{arg}
    default:
      return nil
    }
  }
}

func runValue(_ e: () -> SValue) -> () {
  print(e())
}


struct SName {
  let get: () -> SValue
  var val: SValue {
    return get()
  }
  init(_ v: @escaping () -> SValue) {
    get = v
  }
}

extension SName : EDSL {
  init(integerLiteral : Int) {
    self = SName{.Int(integerLiteral)}
  }
  static func add(_ lhs: SName, _ rhs: SName) -> SName? {
    return SName{SValue.add(lhs.val, rhs.val)!}
  }
  static func mul(_ lhs: SName, _ rhs: SName) -> SName? {
    return SName{SValue.mul(lhs.val, rhs.val)!}
  }
  static func lam(_ body: @escaping (SName) -> SName) -> SName {
    return SName{.Lam { x in body(SName(x)).val }}
  }
  static func app(_ fn: SName, _ arg: SName) -> SName? {
    switch (fn.val) {
    case .Lam(let f):
      return SName{f{arg.val}}
    default:
      return nil
    }
  }
}

func runName(_ e: () -> SName) -> () {
  print(e().get())
}


struct SLazy {
  let get: () -> SValue
  var val: SValue {
    return get()
  }
  init(_ v: @escaping () -> SValue) {
    var cache:SValue? = nil
    get = {
      if cache == nil {
        cache = v()
      }
      return cache!
    }
  }
}

extension SLazy : EDSL {
  init(integerLiteral : Int) {
    self = SLazy{.Int(integerLiteral)}
  }
  static func add(_ lhs: SLazy, _ rhs: SLazy) -> SLazy? {
    return SLazy{SValue.add(lhs.val, rhs.val)!}
  }
  static func mul(_ lhs: SLazy, _ rhs: SLazy) -> SLazy? {
    return SLazy{SValue.mul(lhs.val, rhs.val)!}
  }
  static func lam(_ body: @escaping (SLazy) -> SLazy) -> SLazy {
    return SLazy{.Lam { x in body(SLazy(x)).val }}
  }
  static func app(_ fn: SLazy, _ arg: SLazy) -> SLazy? {
    switch (fn.val) {
    case .Lam(let f):
      return SLazy{f{arg.val}}
    default:
      return nil
    }
  }
}

func runLazy(_ e: () -> SLazy) -> () {
  print(e().get())
}



func test<Exp: EDSL>() -> Exp {
  return
    Exp.let_(2 * 3) { z in
    Exp.let_(5 + 5) { x in
    Exp.let_(x + x) { y in
      y + y
    }}}
}

print("Call by value:")
runValue(test)

print("Call by name:")
runName(test)

print("Lazy evaluation:")
runLazy(test)