Generate curry functions for swift, written in swift

currygen.swift

#!/usr/bin/env swift

import Foundation

infix operator <*> {
    associativity left
    precedence 120
}

func <*> <T: CollectionType, U: CollectionType, V where T.Generator.Element == ((U.Generator.Element) -> V)>(left: T, right: U) -> [V] {
    return left.flatMap { right.map($0) }
}

func <*> <T: CollectionType, U: CollectionType, V where U.Generator.Element == ((T.Generator.Element) -> V)>(left: T, right: U) -> [V] {
    return left.flatMap { right <*> [$0] }
}

extension Array {
    var tail: ArraySlice<Element> {
        return self[1..<count]
    }
    
    var head: ArraySlice<Element> {
        return self[0..<count-1]
    }
}

extension ArraySlice {
    var tail: ArraySlice<Element> {
        return self[startIndex.advancedBy(1)..<endIndex]
    }
    
    var head: ArraySlice<Element> {
        return self[startIndex..<endIndex.advancedBy(-1)]
    }
}

extension Int {
    var tabs: String {
        return Array(count: self, repeatedValue: "    ").joinWithSeparator("")
    }
}

extension UnicodeScalar: ForwardIndexType {
    public func successor() -> UnicodeScalar {
        return UnicodeScalar(value.successor())
    }
}

func recurse<T, U>(f: (T, (T) -> U) -> U) -> (T) -> U {
    var closure: ((T) -> U)!
    closure = { f($0, closure) }
    return closure
}

func identity<T>(value: T) -> T {
    return value
}

func curry<A, B, C>(function: (A, B) -> C) -> (A) -> ((B) -> C) {
    return { a in
        return { b in
            return function(a, b)
        }
    }
}

func log26(value: Double) -> Double {
    return log(value) / log(26)
}

func generateCurryFor<T: SequenceType where T.Generator.Element == String>(variablesSequence: T) -> String {
    // We want an array of strings from this sequence.
    let variables = variablesSequence.map(identity)
    guard variables.count > 2 else { return "" }
    
    // Some cached values used below in multiple cases - the variables comma separated with and without the last value.
    let last = variables[variables.count - 1]
    let secondToLast = variables[variables.count - 2]
    let commaJoinedVariablesHeadHead = variables.head.head.joinWithSeparator(", ")
    let commaJoinedVariablesHead = commaJoinedVariablesHeadHead + ", " + secondToLast
    let commaJoinedVariables = commaJoinedVariablesHead + ", " + last
    
    let decl = "func curry<\(commaJoinedVariables)>(function: (\(commaJoinedVariablesHead)) -> \(last)) -> "
        + recurse({ (strings, me) in
            guard let first = strings.first else { return "" }
            if strings.count == 1 {
                return first
            }
            switch strings.count {
            case 1:
                return first
            case 2:
                return "(\(first)) -> \(strings.last!)"
            default:
                return "(\(first)) -> (\(me(strings.tail)))"
            }
        })(ArraySlice(variables))
    let body = "\(1.tabs)return \(variables.count > 3 ? "curry" : "") { \(commaJoinedVariablesHeadHead.lowercaseString) in\n"
        + "\(2.tabs)return { \(secondToLast.lowercaseString) in\n"
        + "\(3.tabs)return function(\(commaJoinedVariablesHead.lowercaseString))\n"
        + "\(2.tabs)}\n"
        + "\(1.tabs)}"
    
    return decl + " {\n" + body + "\n}"
}

func generateVariables(count: Int) -> [String] {
    guard count > 1 else { return Array(count: Int(count), repeatedValue: "A") }
    let nLetters = Int(ceil(log26(Double(count))))
    let alphabet = Array("A"..."Z" as Range<UnicodeScalar>).map { String($0) }
    return Array((recurse { (index: Int, me: (Int) -> [String]) -> [String] in
        switch index {
        case let x where x <= 0:
            return []
        case 1:
            return alphabet
        default:
            return alphabet + (me(index - 1).map { curry(+)($0) } <*> alphabet)
        }
    })(nLetters)[0..<count])
}

func generateCurriesFor(numberOfVariables number: Int) -> [String] {
    guard number >= 3 else { return [] }
    return (3...number).map { generateCurryFor(generateVariables($0)) }
}

func printUsage() {
    print("Usage: swiftcurrygen [numberOfCurries]")
}

guard Process.argc == 2 else {
    printUsage()
    exit(1)
}

guard let numberOfCurries = Int(Process.arguments[1]) where numberOfCurries > 0 else {
    printUsage()
    exit(1)
}

print(generateCurriesFor(numberOfVariables: numberOfCurries + 2).joinWithSeparator("\n\n"))

// vi:syntax=swift:softtabstop=4:tabstop=4:shiftwidth=4