Refinement.swift

indirect enum Type : CustomStringConvertible {
  case nat
  case product(Type, Type)
  case arrow(Type, Type)

  var description : String {
    switch self {
    case .nat:
      return "ℕ"
    case let .product(t1, t2):
      return "(\(t1) , \(t2))"
    case let .arrow(.arrow(ss, sr), t2):
      return "(\(ss) -> \(sr)) -> \(t2)"
    case let .arrow(t1, t2):
      return "\(t1) -> \(t2)"
    }
  }
}

indirect enum Expression : CustomStringConvertible {
  case evar(String)
  case pair(Expression, Expression)
  case fst(Type, Expression)
  case snd(Type, Expression)
  case abs(String, Expression)
  case app(Type, Expression, Expression)

  var description : String {
    switch self {
    case let .evar(name):
      return name
    case let .app(_, e1, e2):
      switch e2 {
      case .abs(_, _):
        return "\(e1)(\(e2))"
      case .app(_, _, _):
        return "\(e1) (\(e2))"
      default:
        return "\(e1) \(e2)"
      }
    case let .abs(n, e):
      return "λ \(n) -> \(e)"
    case let .pair(a, b):
      return "(\(a) , \(b))"
    case let .fst(_, e):
      return "fst(\(e))"
    case let .snd(_, e):
      return "snd(\(e))"
    }
  }
}

protocol Problem : CustomStringConvertible {
  var decompose : [Self]? { get }
}

enum Judgement: Problem {
  case equal(Type, Type)
  case inhabited(Dictionary<String, Type>, Expression, Type)

  var description : String {
    switch self {
    case let .equal(ty1, ty2):
      return "\(ty1) == \(ty2)"
    case let .inhabited(_, e, ty):
      return "\(e) : \(ty)"
    }
  }

  var decompose : [Judgement]? {
    switch self {
    case let .equal(type1, type2):
      switch (type1, type2) {
      case (.nat, .nat):
        return .some([])
      case let (.product(a1, b1), .product(a2, b2)):
        return .some([.equal(a1, a2), .equal(b1, b2)])
      case let (.arrow(a1, b1), .arrow(a2, b2)):
        return .some([.equal(a1, a2), .equal(b1, b2)])
      default:
        return .none
      }
    case let .inhabited(gamma, expr, type):
      switch (expr, type) {
      case let (.evar(x), a):
        return gamma[x].map { [.equal(a, $0)] }
      case let (.pair(m, n), .product(a, b)):
        return .some([ .inhabited(gamma, m, a), .inhabited(gamma, n, b) ])
      case let (.fst(b, p), a):
        return .some([ .inhabited(gamma, p, .product(a, b)) ])
      case let (.snd(a, p), b):
        return .some([ .inhabited(gamma, p, .product(a, b)) ])
      case let (.abs(x, m), .arrow(a, b)):
        var g2 = gamma
        g2[x] = a
        return .some([ .inhabited(g2, m, b), ])
      case let (.app(a, m, n), b):
        return .some([ .inhabited(gamma, m, .arrow(a, b)), .inhabited(gamma, n, a) ])
      default:
        return .none
      }
    }
  }
}

indirect enum ProofTree<P : Problem> /*: CustomStringConvertible*/ {
  case branch(P, [ProofTree])
}

struct Prob<Lem : Problem> {
  let lem : Lem
  init(lem : Lem) { self.lem = lem }

  func solve() -> ProofTree<Lem>? {
    return self.lem.decompose.flatMap { (js : [Lem]) -> ProofTree<Lem>? in
      return sequence(js.map({ Prob(lem: $0).solve() })).map { ts in ProofTree.branch(self.lem, ts) }
    }
  }
}

private func sequence<A>(_ ms : [Optional<A>]) -> Optional<[A]> {
  return ms.reduce(Optional<[A]>.some([]), { n, m in
    return m.flatMap { x in
      return n.flatMap { xs in
        return Optional<[A]>.some(xs + [x])
      }
    }
  })
}

let d = Prob(lem:
  Judgement.inhabited(["x" : .nat, "f" : .arrow(.nat, .nat)], .app(.nat, .evar("f"), .app(.nat, .evar("f"), .evar("x"))), .nat)
).solve()!