pidigits benchmark in Swift

Include_swift_gmp_module.modulemap

module gmp [system] {
    header "shim.h"
    link "gmp"
    export *
}

Include_swift_gmp_shim.h

#include <gmp.h>

pidigits.swift

// Implementation of the pidigits benchmarks game benchmark
// https://benchmarksgame-team.pages.debian.net/benchmarksgame/performance/pidigits.html
//
// ~/Library/Developer/Toolchains/swift-DEVELOPMENT-SNAPSHOT-2026-04-28-a.xctoolchain/usr/bin/swiftc \
//    -sdk "$(xcrun --show-sdk-path)" \                       
//    pidigits.swift \                                                                                                                   
//    -O -wmo -parse-as-library \                                                                                                        
//    -enable-experimental-feature Lifetimes \                                                                                           
//    -I Include/swift/gmp \                                                                                                             
//    -Xcc -I/opt/homebrew/include \                          
//    -L/opt/homebrew/lib \                                                                                                              
//    -o pidigits

import gmp

struct BigInt: ~Copyable {
    var storage: mpz_t = mpz_t()
    // Per-instance scratch so in-place ops can write into a separate buffer
    // and swap, avoiding the GMP rop==op aliasing requirement (which would
    // otherwise force a `withUnsafeMutablePointer` closure to bind `&storage`
    // once and pass it twice).
    var scratch: mpz_t = mpz_t()

    init() {
        __gmpz_init(&storage)
        __gmpz_init(&scratch)
    }

    init(_ value: UInt) {
        __gmpz_init_set_ui(&storage, value)
        __gmpz_init(&scratch)
    }

    deinit {
        var s = storage
        __gmpz_clear(&s)
        var sc = scratch
        __gmpz_clear(&sc)
    }

    // Lightweight (mpz, scalar) pair produced by `BigInt * UInt`. Consumed by
    // `+=` (→ __gmpz_addmul_ui), `-=` (→ __gmpz_submul_ui), and `<-`
    // (→ __gmpz_mul_ui). Aliases the source's _mp_d limb buffer, so it must
    // not outlive the source — enforced by `~Escapable` + `@_lifetime`.
    struct Scaled: ~Copyable, ~Escapable {
        let storage: mpz_t
        let scale: UInt

        @_lifetime(borrow source)
        init(_ source: borrowing BigInt, scale: UInt) {
            self.storage = source.storage
            self.scale = scale
        }
    }

    // Deferred binary op produced by `a + b` and `a / b`, consumed by `<-`.
    struct Sum: ~Copyable, ~Escapable {
        let a: mpz_t
        let b: mpz_t

        @_lifetime(borrow first, borrow second)
        init(_ first: borrowing BigInt, _ second: borrowing BigInt) {
            self.a = first.storage
            self.b = second.storage
        }
    }
    struct Quotient: ~Copyable, ~Escapable {
        let a: mpz_t
        let b: mpz_t

        @_lifetime(borrow first, borrow second)
        init(_ first: borrowing BigInt, _ second: borrowing BigInt) {
            self.a = first.storage
            self.b = second.storage
        }
    }

    @_lifetime(borrow lhs)
    static func * (lhs: borrowing BigInt, rhs: UInt) -> Scaled {
        Scaled(lhs, scale: rhs)
    }

    @_lifetime(borrow lhs, borrow rhs)
    static func + (lhs: borrowing BigInt, rhs: borrowing BigInt) -> Sum {
        Sum(lhs, rhs)
    }

    @_lifetime(borrow lhs, borrow rhs)
    static func / (lhs: borrowing BigInt, rhs: borrowing BigInt) -> Quotient {
        Quotient(lhs, rhs)
    }

    static func *= (lhs: inout BigInt, rhs: UInt) {
        __gmpz_mul_ui(&lhs.scratch, &lhs.storage, rhs)
        swap(&lhs.storage, &lhs.scratch)
    }

    static func += (lhs: inout BigInt, rhs: consuming Scaled) {
        var src = rhs.storage
        __gmpz_addmul_ui(&lhs.storage, &src, rhs.scale)
    }

    static func += (lhs: inout BigInt, rhs: borrowing BigInt) {
        var src = rhs.storage
        __gmpz_add(&lhs.scratch, &lhs.storage, &src)
        swap(&lhs.storage, &lhs.scratch)
    }

    static func -= (lhs: inout BigInt, rhs: consuming Scaled) {
        var src = rhs.storage
        __gmpz_submul_ui(&lhs.storage, &src, rhs.scale)
    }

    static func <- (lhs: inout BigInt, rhs: consuming Scaled) {
        var src = rhs.storage
        __gmpz_mul_ui(&lhs.storage, &src, rhs.scale)
    }

    static func <- (lhs: inout BigInt, rhs: consuming Sum) {
        var a = rhs.a
        var b = rhs.b
        __gmpz_add(&lhs.storage, &a, &b)
    }

    static func <- (lhs: inout BigInt, rhs: consuming Quotient) {
        var a = rhs.a
        var b = rhs.b
        __gmpz_tdiv_q(&lhs.storage, &a, &b)
    }
}

infix operator <- : AssignmentPrecedence

extension BigInt: Comparable {
    static func == (lhs: borrowing BigInt, rhs: borrowing BigInt) -> Bool {
        var l = lhs.storage
        var r = rhs.storage
        return __gmpz_cmp(&l, &r) == 0
    }

    static func < (lhs: borrowing BigInt, rhs: borrowing BigInt) -> Bool {
        var l = lhs.storage
        var r = rhs.storage
        return __gmpz_cmp(&l, &r) < 0
    }
}

extension UInt8 {
    init(_ mpz: borrowing BigInt) {
        var s = mpz.storage
        self = UInt8(__gmpz_get_ui(&s))
    }
}

@main
struct PiDigits: ~Copyable {
    var tmp1 = BigInt()
    var tmp2 = BigInt()
    var acc = BigInt(0)
    var den = BigInt(1)
    var num = BigInt(1)
    var k: UInt = 0

    mutating func nextTerm(_ k: UInt) {
        let k2 = 2 * k + 1
        acc += num * 2
        acc *= k2
        den *= k2
        num *= k
    }

    // Computes one digit if it has converged. Returns nil if the algorithm needs
    // more inner iterations, or the digit value if d == d4.
    //
    // d  = floor((num*3 + acc) / den)
    // d4 = floor((num*4 + acc) / den)
    // num*4 + acc = (num*3 + acc) + num, so we reuse the intermediate to save
    // one mul_ui per outer iteration.
    mutating func extractDigit() -> UInt8? {
        tmp1 <- num * 3
        tmp2 <- tmp1 + acc                     // tmp2 = num*3 + acc
        tmp1 <- tmp2 / den                     // tmp1 = d
        let d = UInt8(tmp1)
        tmp2 += num                            // tmp2 = num*4 + acc
        tmp1 <- tmp2 / den                     // tmp1 = d4
        return d == UInt8(tmp1) ? d : nil
    }

    mutating func eliminateDigit(_ d: UInt8) {
        acc -= den * UInt(d)
        acc *= 10
        num *= 10
    }

    mutating func nextDigit() -> UInt8 {
        while true {
            repeat {
                k += 1
                nextTerm(k)
            } while num > acc

            if let d = extractDigit() {
                eliminateDigit(d)
                return d
            }
        }
    }

    static func main() {
        let n = CommandLine.arguments.last.flatMap(Int.init) ?? 50

        var digits: [UInt8] = []
        digits.reserveCapacity(n)
        var pi = PiDigits()
        for _ in 0..<n {
            digits.append(pi.nextDigit())
        }

        for chunk in stride(from: 0, to: n, by: 10) {
            let count = min(10, n - chunk)
            let digitsStr = digits[chunk..<chunk + count].map(String.init).joined()
            let padding = String(repeating: " ", count: 10 - count)
            print("\(digitsStr)\(padding)\t:\(chunk + count)")
        }
    }
}