Szudzik's Elegant Pairing Function in haskell

Pairing.hs

-- | Szudzik's Elegant Pairing Function
--
-- http://szudzik.com/ElegantPairing.pdf
--
-- For all non-negative integers:
--
-- @
-- uncurry pair . unpair = id
-- unpair . uncurry pair = id
-- @
module Pairing (pair, unpair) where


-- | Pack two integers into one.
pair :: Integer -> Integer -> Integer
pair y x =
  if y > x then
    y * y + x
  else
    x * x + x + y


-- | Unpack one integer into two.
unpair :: Integer -> (Integer, Integer)
unpair z =
  if l < q then
    (q, l)
  else
    (l-q, q)
  where
    q = squareRoot z
    l = z - q ^ two


-- Adapted from https://wiki.haskell.org/Generic_number_type
squareRoot :: Integer -> Integer
squareRoot 0 = 0
squareRoot 1 = 1
squareRoot n =
   let twopows = iterate (^ two) two
       (lowerRoot, lowerN) =
          last $ takeWhile ((n>=) . snd) $ zip (1:twopows) twopows
       newtonStep x = div (x + div n x) two
       iters = iterate newtonStep (squareRoot (div n lowerN) * lowerRoot)
       isRoot r  =  r ^ two <= n && n < (r+1) ^ two
   in  head $ dropWhile (not . isRoot) iters


-- Defeat type defaulting without specifying the type every time.
two :: Integer
two = 2
{-# INLINE two #-}