Proof-of-concept Haskell implementation of LCG fast-forwarding

FastForwardLCG.hs

-- An LCG consists of three values: a multiplier, an offset, and a
-- modulus.
data LCG = LCG
  { multiplier :: Integer
  , offset     :: Integer
  , modulus    :: Integer
  }
  deriving Show

-- Use an LCG to advance from one value to the next.  LCG a c n
-- represents the function which takes a current seed x and maps it to
-- ax+c (mod n).
step :: LCG -> Integer -> Integer
step (LCG a c n) = \x -> (a*x + c) `mod` n

-- l1 <.> l2 is the *composition* of l1 and l2; that is, it yields the
-- LCG which performs a step of l2 followed by a step of l1.
-- (Precondition: l1 and l2 have the same modulus.)
--
-- If we first apply l2 to x, we get a2*x + c2; then applying l1 to this
-- result in turn yields
--
--   a1*(a2*x + c2) + c1 = a1*a2*x + a1*c2 + c1.
--
-- Hence the composition can be represented by the LCG with multiplier
-- a1*a2 and offset a1*c2 + c1.
(<.>) :: LCG -> LCG -> LCG
(LCG a1 c1 n) <.> (LCG a2 c2 _) = LCG ((a1*a2) `mod` n) ((a1*c2 + c1) `mod` n) n

infixr 9 <.>

-- Raise an LCG to an "exponent"; that is, (powerLCG l k) yields the LCG
-- which is equal to l composed with itself k times.  Uses a repeated
-- squaring algorithm which takes at most about 2 log_2(k) composition
-- operations.
powerLCG :: LCG -> Integer -> LCG
powerLCG lcg 1 = lcg
powerLCG lcg n

  -- To compute lcg^n, we compute lcg^(n/2), and then square it
  -- (i.e. compose with itself) if n is even, or else square it and
  -- compose with one more copy of lcg if n is odd.
  | even n    =         half <.> half
  | otherwise = lcg <.> half <.> half
  where
    half = powerLCG lcg (n `div` 2)   -- n `div` 2 rounds down

-- We can do some quick sanity checks to make sure this gives us the
-- right thing.  For example:

-- Get the 943rd value generated by ran0 starting with the initial
-- seed 666.  This version works by iterating ran0 step-by-step and
-- getting the 943rd element of the resulting list.
value1 :: Integer
value1 = iterate (step ran0) 666 !! 943

-- On the other hand, we could first raise ran0 to the 943rd power and
-- then apply the result to 666.  This is MUCH faster.
value2 :: Integer
value2 = step (powerLCG ran0 943) 666

-- As you can check, value1 == value2 == 1707103193 .

------------------------------------------------------------

initseed :: Integer
initseed = 666

ia, im :: Integer
ia = 16807
im = 2^31 - 1

ran0 :: LCG
ran0 = LCG ia 0 im