-- | Szudzik's Elegant Pairing Function
--
-- http://szudzik.com/ElegantPairing.pdf
--
-- For all non-negative integers:
--
-- @
-- uncurry pair . unpair = id
-- unpair . uncurry pair = id
-- @
-- Addapted from https://gist.github.com/klntsky/7026018c3341e6aa17bc237746ee0256#file-pairing-hs
-- We use Int rather than Integer, and need accuracy in the range [0..10000]
module Test.Cardano.Ledger.Constrained.Pairing (pair, unpair) where

-- | Pack two integers into one.
pair :: Int -> Int -> Int
pair :: Int -> Int -> Int
pair Int
y Int
x =
  if Int
y Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
x
    then Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x
    else Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y

-- | Unpack one integer into two.
unpair :: Int -> (Int, Int)
unpair :: Int -> (Int, Int)
unpair Int
z =
  if Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
q
    then (Int
q, Int
l)
    else (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
q, Int
q)
  where
    q :: Int
q = Int -> Int
squareRoot Int
z
    l :: Int
l = Int
z Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
q Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
two

-- Adapted from https://wiki.haskell.org/Generic_number_type
squareRoot :: Int -> Int
squareRoot :: Int -> Int
squareRoot Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> Int
forall a. HasCallStack => [Char] -> a
error ([Char]
"squareRoot works only for positive Int: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n)
squareRoot Int
0 = Int
0
squareRoot Int
1 = Int
1
squareRoot Int
n =
  let twopows :: [Int]
twopows = (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
two) Int
two
      (Int
lowerRoot, Int
lowerN) =
        [(Int, Int)] -> (Int, Int)
forall a. HasCallStack => [a] -> a
last ([(Int, Int)] -> (Int, Int)) -> [(Int, Int)] -> (Int, Int)
forall a b. (a -> b) -> a -> b
$ ((Int, Int) -> Bool) -> [(Int, Int)] -> [(Int, Int)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>=) (Int -> Bool) -> ((Int, Int) -> Int) -> (Int, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> Int
forall a b. (a, b) -> b
snd) ([(Int, Int)] -> [(Int, Int)]) -> [(Int, Int)] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Int
1 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
twopows) [Int]
twopows
      newtonStep :: Int -> Int
newtonStep Int
x = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
n Int
x) Int
two
      iters :: [Int]
iters = (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate Int -> Int
newtonStep (Int -> Int
squareRoot (Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
n Int
lowerN) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
lowerRoot)
      isRoot :: Int -> Bool
isRoot Int
r = Int
r Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
two Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
two
   in case (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Bool -> Bool
not (Bool -> Bool) -> (Int -> Bool) -> Int -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Bool
isRoot) [Int]
iters of
        (Int
r : [Int]
_) -> Int
r
        [] -> [Char] -> Int
forall a. HasCallStack => [Char] -> a
error ([Char]
"Not possible, every positive Int has an Integral square root: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n)

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