Question

I’m writing a powMod function which I have to use quite intensively. The starting point is a custom pow function:

// Compute power using multiplication and square.
// pow (*) (^2) 1 x n = x^n
let pow mul sq one x n =   
    let rec loop x' n' acc =
       match n' with
       | 0 -> acc
       | _ -> let q = n'/2
              let r = n'%2
              let x2 = sq x'
              if r = 0 then
                 loop x2 q acc
              else
                 loop x2 q (mul x' acc)
    loop x n one

After inspecting the range of my input, I chose int64 because it is big enough to represent output and I can avoid expensive calculation with bigint:

let mulMod m a b = (a*b)%m
let squareMod m a = mulMod m a a
let powMod m = pow (mulMod m) (squareMod m) 1L

I assume that modulo (m) is bigger than multipliers (a, b) and functions work on non-negative numbers only. The powMod function is correct for most cases; however, the problem lies in the mulMod function where a*b may be over int64 range but (a*b)%m is not.
The below example demonstrates the overflow issue:

let a = (pown 2L 40) - 1L
let b = (pown 2L 32) - 1L
let p = powMod a b 2 // p = -8589934591L -- wrong

Is there any way to avoid int64 overflow without resorting to bigint type?

Answer 1

The problem that you have is that all of your intermediate calculations are implicitly mod 2⁶⁴, and it’s not generally true that

a·b mod m = (a·b mod 2⁶⁴) mod m

which is what you’re calculating.

I can’t think of a simple way to do the correct calculation using just 64-bit numbers, but you don’t have to go all the way up to bigints; if a and b have at most 64 bits, then their full product has at most 128 bits, so you can keep track of the product in two 64-bit integers (here bundled as a custom struct):

// bit width of a uint64, needed for mod calculation
let width = 
    let rec loop w = function
    | 0uL -> w
    | n -> loop (w+1) (n >>> 1)
    loop 0

[<Struct; CustomComparison; CustomEquality>]
type UInt128 =
    val hi : uint64
    val lo : uint64
    new (hi,lo) = { lo = lo; hi = hi }
    new (lo) = { lo = lo; hi = 0uL }
    static member (+)(x:UInt128, y:UInt128) =
        if x.lo > 0xffffffffuL - y.lo then
            UInt128(x.hi + y.hi + 1uL, x.lo + y.lo)
        else
            UInt128(x.hi + y.hi, x.lo + y.lo)
    static member (-)(x:UInt128, y:UInt128) =
        if y.lo > x.lo then
            UInt128(x.hi - y.hi - 1uL, x.lo - y.lo)
        else
            UInt128(x.hi - y.hi, x.lo - y.lo)

    static member ( * )(x:UInt128, y:UInt128) =
        let a1 = ((x.lo &&& 0xffffffffuL) * (y.lo &&& 0xffffffffuL)) >>> 32
        let a2 =  (x.lo &&& 0xffffffffuL) * (y.lo >>> 32)
        let a3 =  (x.lo >>> 32) * (y.lo &&& 0xffffffffuL)
        let sum = ((a1 + a2 + a3) >>> 32) + (x.lo >>> 32) * (y.lo >>> 32)
        let sum =
            if a2 > 0xffffffffffffffffuL - a1 || a1 + a2 > 0xffffffffffffffffuL - a3 then
                0x100000000uL + sum
            else
                sum
        UInt128(x.hi * y.lo + x.lo * y.hi + sum, x.lo * y.lo)

    static member (>>>)(x:UInt128, n) =
        UInt128(x.hi >>> n, x.lo >>> n)

    static member (<<<)(x:UInt128, n) =
        UInt128((x.hi <<< n) + (x.lo >>> (64 - n)), x.lo <<< n)

    interface System.IComparable with
        member x.CompareTo(y) =
            match y with
            | :? UInt128 as y ->
                match x.hi.CompareTo(y.hi) with
                | 0 -> x.lo.CompareTo(y.lo)
                | n -> n

    override x.Equals(y) = 
        match y with
        | :? UInt128 as y -> x.hi = y.hi && x.lo = y.lo
        | _ -> false

    override x.GetHashCode() = x.hi.GetHashCode() + x.lo.GetHashCode() * 7

    (* calculate mod via long-division *)
    static member (%)(x:UInt128, d) =
        let rec reduce (r:UInt128) d' =
            if r.hi = 0uL then r.lo % d
            else
                let r' = if r < d' then r else r - d'
                reduce r' (d' >>> 1)
        let shift = width x.hi + (64 - width d)
        reduce x (UInt128(0uL,d) <<< shift)

let mulMod m a b =
    UInt128(a) * UInt128(b) % m

(* squareMod, powMod basically as before: *)
let squareMod m a = mulMod m a a  
let powMod m = pow (mulMod m) (squareMod m) 1uL  

let a = (pown 2uL 40) - 1uL  
let b = (pown 2uL 32) - 1uL  
let p = powMod a b 2

Having said that, since bigints will give you the correct answer, why not just use bigints to do the intermediate calculation and convert to long at the end (which is guaranteed to be a lossless conversion given m’s range)? I suspect that the performance penalty for using bigints should be acceptable for most applications (compared to the headache of maintaining your own math routines).