-- |
-- Module      :  Cryptol.TypeCheck.Solver.InfNat
-- Copyright   :  (c) 2013-2016 Galois, Inc.
-- License     :  BSD3
-- Maintainer  :  cryptol@galois.com
-- Stability   :  provisional
-- Portability :  portable
--
-- This module defines natural numbers with an additional infinity
-- element, and various arithmetic operators on them.

{-# LANGUAGE Safe #-}

{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
module Cryptol.TypeCheck.Solver.InfNat where

import Data.Bits
import Cryptol.Utils.Panic

import GHC.Generics (Generic)
import Control.DeepSeq

-- | Natural numbers with an infinity element
data Nat' = Nat Integer | Inf
            deriving (Int -> Nat' -> ShowS
[Nat'] -> ShowS
Nat' -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Nat'] -> ShowS
$cshowList :: [Nat'] -> ShowS
show :: Nat' -> String
$cshow :: Nat' -> String
showsPrec :: Int -> Nat' -> ShowS
$cshowsPrec :: Int -> Nat' -> ShowS
Show, Nat' -> Nat' -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Nat' -> Nat' -> Bool
$c/= :: Nat' -> Nat' -> Bool
== :: Nat' -> Nat' -> Bool
$c== :: Nat' -> Nat' -> Bool
Eq, Eq Nat'
Nat' -> Nat' -> Bool
Nat' -> Nat' -> Ordering
Nat' -> Nat' -> Nat'
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Nat' -> Nat' -> Nat'
$cmin :: Nat' -> Nat' -> Nat'
max :: Nat' -> Nat' -> Nat'
$cmax :: Nat' -> Nat' -> Nat'
>= :: Nat' -> Nat' -> Bool
$c>= :: Nat' -> Nat' -> Bool
> :: Nat' -> Nat' -> Bool
$c> :: Nat' -> Nat' -> Bool
<= :: Nat' -> Nat' -> Bool
$c<= :: Nat' -> Nat' -> Bool
< :: Nat' -> Nat' -> Bool
$c< :: Nat' -> Nat' -> Bool
compare :: Nat' -> Nat' -> Ordering
$ccompare :: Nat' -> Nat' -> Ordering
Ord, forall x. Rep Nat' x -> Nat'
forall x. Nat' -> Rep Nat' x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Nat' x -> Nat'
$cfrom :: forall x. Nat' -> Rep Nat' x
Generic, Nat' -> ()
forall a. (a -> ()) -> NFData a
rnf :: Nat' -> ()
$crnf :: Nat' -> ()
NFData)

fromNat :: Nat' -> Maybe Integer
fromNat :: Nat' -> Maybe Integer
fromNat Nat'
n' =
  case Nat'
n' of
    Nat Integer
i -> forall a. a -> Maybe a
Just Integer
i
    Nat'
_     -> forall a. Maybe a
Nothing




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


nAdd :: Nat' -> Nat' -> Nat'
nAdd :: Nat' -> Nat' -> Nat'
nAdd Nat'
Inf Nat'
_           = Nat'
Inf
nAdd Nat'
_ Nat'
Inf           = Nat'
Inf
nAdd (Nat Integer
x) (Nat Integer
y) = Integer -> Nat'
Nat (Integer
x forall a. Num a => a -> a -> a
+ Integer
y)

{-| Some algebraic properties of interest:

> 1 * x = x
> x * (y * z) = (x * y) * z
> 0 * x = 0
> x * y = y * x
> x * (a + b) = x * a + x * b

-}
nMul :: Nat' -> Nat' -> Nat'
nMul :: Nat' -> Nat' -> Nat'
nMul (Nat Integer
0) Nat'
_       = Integer -> Nat'
Nat Integer
0
nMul Nat'
_ (Nat Integer
0)       = Integer -> Nat'
Nat Integer
0
nMul Nat'
Inf Nat'
_           = Nat'
Inf
nMul Nat'
_ Nat'
Inf           = Nat'
Inf
nMul (Nat Integer
x) (Nat Integer
y) = Integer -> Nat'
Nat (Integer
x forall a. Num a => a -> a -> a
* Integer
y)


{-| Some algebraic properties of interest:

> x ^ 0        = 1
> x ^ (n + 1)  = x * (x ^ n)
> x ^ (m + n)  = (x ^ m) * (x ^ n)
> x ^ (m * n)  = (x ^ m) ^ n

-}
nExp :: Nat' -> Nat' -> Nat'
nExp :: Nat' -> Nat' -> Nat'
nExp Nat'
_ (Nat Integer
0)       = Integer -> Nat'
Nat Integer
1
nExp Nat'
Inf Nat'
_           = Nat'
Inf
nExp (Nat Integer
0) Nat'
Inf     = Integer -> Nat'
Nat Integer
0
nExp (Nat Integer
1) Nat'
Inf     = Integer -> Nat'
Nat Integer
1
nExp (Nat Integer
_) Nat'
Inf     = Nat'
Inf
nExp (Nat Integer
x) (Nat Integer
y) = Integer -> Nat'
Nat (Integer
x forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
y)

nMin :: Nat' -> Nat' -> Nat'
nMin :: Nat' -> Nat' -> Nat'
nMin Nat'
Inf Nat'
x            = Nat'
x
nMin Nat'
x Nat'
Inf            = Nat'
x
nMin (Nat Integer
x) (Nat Integer
y)  = Integer -> Nat'
Nat (forall a. Ord a => a -> a -> a
min Integer
x Integer
y)

nMax :: Nat' -> Nat' -> Nat'
nMax :: Nat' -> Nat' -> Nat'
nMax Nat'
Inf Nat'
_            = Nat'
Inf
nMax Nat'
_ Nat'
Inf            = Nat'
Inf
nMax (Nat Integer
x) (Nat Integer
y)  = Integer -> Nat'
Nat (forall a. Ord a => a -> a -> a
max Integer
x Integer
y)

{- | @nSub x y = Just z@ iff @z@ is the unique value
such that @Add y z = Just x@.  -}
nSub :: Nat' -> Nat' -> Maybe Nat'
nSub :: Nat' -> Nat' -> Maybe Nat'
nSub Nat'
Inf (Nat Integer
_)              = forall a. a -> Maybe a
Just Nat'
Inf
nSub (Nat Integer
x) (Nat Integer
y)
  | Integer
x forall a. Ord a => a -> a -> Bool
>= Integer
y                    = forall a. a -> Maybe a
Just (Integer -> Nat'
Nat (Integer
x forall a. Num a => a -> a -> a
- Integer
y))
nSub Nat'
_ Nat'
_                      = forall a. Maybe a
Nothing


-- XXX:
-- Does it make sense to define:
--   nDiv Inf (Nat x)  = Inf
--   nMod Inf (Nat x)  = Nat 0

{- | Rounds down.

> y * q + r = x
> x / y     = q with remainder r
> 0 <= r && r < y

We don't allow `Inf` in the first argument for two reasons:
  1. It matches the behavior of `nMod`,
  2. The well-formedness constraints can be expressed as a conjunction.
-}
nDiv :: Nat' -> Nat' -> Maybe Nat'
nDiv :: Nat' -> Nat' -> Maybe Nat'
nDiv Nat'
_       (Nat Integer
0)  = forall a. Maybe a
Nothing
nDiv Nat'
Inf Nat'
_            = forall a. Maybe a
Nothing
nDiv (Nat Integer
x) (Nat Integer
y)  = forall a. a -> Maybe a
Just (Integer -> Nat'
Nat (forall a. Integral a => a -> a -> a
div Integer
x Integer
y))
nDiv (Nat Integer
_) Nat'
Inf      = forall a. a -> Maybe a
Just (Integer -> Nat'
Nat Integer
0)

nMod :: Nat' -> Nat' -> Maybe Nat'
nMod :: Nat' -> Nat' -> Maybe Nat'
nMod Nat'
_       (Nat Integer
0)  = forall a. Maybe a
Nothing
nMod Nat'
Inf     Nat'
_        = forall a. Maybe a
Nothing
nMod (Nat Integer
x) (Nat Integer
y)  = forall a. a -> Maybe a
Just (Integer -> Nat'
Nat (forall a. Integral a => a -> a -> a
mod Integer
x Integer
y))
nMod (Nat Integer
x) Nat'
Inf      = forall a. a -> Maybe a
Just (Integer -> Nat'
Nat Integer
x)          -- inf * 0 + x = 0 + x

-- | @nCeilDiv msgLen blockSize@ computes the least @n@ such that
-- @msgLen <= blockSize * n@. It is undefined when @blockSize = 0@,
-- or when @blockSize = inf@. @inf@ divided by any positive
-- finite value is @inf@.
nCeilDiv :: Nat' -> Nat' -> Maybe Nat'
nCeilDiv :: Nat' -> Nat' -> Maybe Nat'
nCeilDiv Nat'
_       Nat'
Inf      = forall a. Maybe a
Nothing
nCeilDiv Nat'
_       (Nat Integer
0)  = forall a. Maybe a
Nothing
nCeilDiv Nat'
Inf     (Nat Integer
_)  = forall a. a -> Maybe a
Just Nat'
Inf
nCeilDiv (Nat Integer
x) (Nat Integer
y)  = forall a. a -> Maybe a
Just (Integer -> Nat'
Nat (- forall a. Integral a => a -> a -> a
div (- Integer
x) Integer
y))

-- | @nCeilMod msgLen blockSize@ computes the least @k@ such that
-- @blockSize@ divides @msgLen + k@. It is undefined when @blockSize = 0@
-- or @blockSize = inf@.  @inf@ modulus any positive finite value is @0@.
nCeilMod :: Nat' -> Nat' -> Maybe Nat'
nCeilMod :: Nat' -> Nat' -> Maybe Nat'
nCeilMod Nat'
_       Nat'
Inf      = forall a. Maybe a
Nothing
nCeilMod Nat'
_       (Nat Integer
0)  = forall a. Maybe a
Nothing
nCeilMod Nat'
Inf     (Nat Integer
_)  = forall a. a -> Maybe a
Just (Integer -> Nat'
Nat Integer
0)
nCeilMod (Nat Integer
x) (Nat Integer
y)  = forall a. a -> Maybe a
Just (Integer -> Nat'
Nat (forall a. Integral a => a -> a -> a
mod (- Integer
x) Integer
y))

-- | Rounds up.
-- @lg2 x = y@, iff @y@ is the smallest number such that @x <= 2 ^ y@
nLg2 :: Nat' -> Nat'
nLg2 :: Nat' -> Nat'
nLg2 Nat'
Inf      = Nat'
Inf
nLg2 (Nat Integer
0)  = Integer -> Nat'
Nat Integer
0
nLg2 (Nat Integer
n)  = case Integer -> Integer -> Maybe (Integer, Bool)
genLog Integer
n Integer
2 of
                  Just (Integer
x,Bool
exact) | Bool
exact     -> Integer -> Nat'
Nat Integer
x
                                 | Bool
otherwise -> Integer -> Nat'
Nat (Integer
x forall a. Num a => a -> a -> a
+ Integer
1)
                  Maybe (Integer, Bool)
Nothing -> forall a. HasCallStack => String -> [String] -> a
panic String
"Cryptol.TypeCheck.Solver.InfNat.nLg2"
                               [ String
"genLog returned Nothing" ]

-- | @nWidth n@ is number of bits needed to represent all numbers
-- from 0 to n, inclusive. @nWidth x = nLg2 (x + 1)@.
nWidth :: Nat' -> Nat'
nWidth :: Nat' -> Nat'
nWidth Nat'
Inf      = Nat'
Inf
nWidth (Nat Integer
n)  = Integer -> Nat'
Nat (Integer -> Integer
widthInteger Integer
n)



-- | @length [ x, y .. z ]@
nLenFromThenTo :: Nat' -> Nat' -> Nat' -> Maybe Nat'
nLenFromThenTo :: Nat' -> Nat' -> Nat' -> Maybe Nat'
nLenFromThenTo (Nat Integer
x) (Nat Integer
y) (Nat Integer
z)
  | Integer
step forall a. Eq a => a -> a -> Bool
/= Integer
0 = let len :: Integer
len = forall a. Integral a => a -> a -> a
div Integer
dist Integer
step forall a. Num a => a -> a -> a
+ Integer
1
                in forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Integer -> Nat'
Nat forall a b. (a -> b) -> a -> b
$ if Integer
x forall a. Ord a => a -> a -> Bool
> Integer
y
                                  -- decreasing
                                  then (if Integer
z forall a. Ord a => a -> a -> Bool
> Integer
x then Integer
0 else Integer
len)
                                  -- increasing
                                  else (if Integer
z forall a. Ord a => a -> a -> Bool
< Integer
x then Integer
0 else Integer
len)
  where
  step :: Integer
step = forall a. Num a => a -> a
abs (Integer
x forall a. Num a => a -> a -> a
- Integer
y)
  dist :: Integer
dist = forall a. Num a => a -> a
abs (Integer
x forall a. Num a => a -> a -> a
- Integer
z)

nLenFromThenTo Nat'
_ Nat'
_ Nat'
_ = forall a. Maybe a
Nothing

{- Note [Sequences of Length 0]

  nLenFromThenTo x y z == 0
    case 1: x > y  && z > x
    case 2: x <= y && z < x
-}

{- Note [Sequences of Length 1]

  `nLenFromThenTo x y z == 1`

  dist < step && (x > y && z <= x   ||   y >= x && z >= x)

  case 1: dist < step && x > y && z <= x
  case 2: dist < step && y >= x && z >= x

  case 1: if    `z <= x`,
          then  `x - z >= 0`,
          hence `dist = x - z`    (a)

          if    `x > y`
          then  `x - y` > 0
          hence `step = x - y`    (b)

          from (a) and (b):
            `dist < step`
            `x - z < x - y`
            `-z < -y`
            `z > y`

  case 1 summary:  x >= z && z > y



  case 2: if y >= x, then step = y - x   (a)
          if z >= x, then dist = z - x   (b)

          dist < step =
          (z - x) < (y - x) =
          (z < y)

  case 2 summary: y > z, z >= x

-}


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

-- | Compute the logarithm of a number in the given base, rounded down to the
-- closest integer.  The boolean indicates if we the result is exact
-- (i.e., True means no rounding happened, False means we rounded down).
-- The logarithm base is the second argument.
genLog :: Integer -> Integer -> Maybe (Integer, Bool)
genLog :: Integer -> Integer -> Maybe (Integer, Bool)
genLog Integer
x Integer
0    = if Integer
x forall a. Eq a => a -> a -> Bool
== Integer
1 then forall a. a -> Maybe a
Just (Integer
0, Bool
True) else forall a. Maybe a
Nothing
genLog Integer
_ Integer
1    = forall a. Maybe a
Nothing
genLog Integer
0 Integer
_    = forall a. Maybe a
Nothing
genLog Integer
x Integer
base = forall a. a -> Maybe a
Just (forall {t}. Num t => t -> Integer -> (t, Bool)
exactLoop Integer
0 Integer
x)
  where
  exactLoop :: t -> Integer -> (t, Bool)
exactLoop t
s Integer
i
    | Integer
i forall a. Eq a => a -> a -> Bool
== Integer
1     = (t
s,Bool
True)
    | Integer
i forall a. Ord a => a -> a -> Bool
< Integer
base   = (t
s,Bool
False)
    | Bool
otherwise  =
        let s1 :: t
s1 = t
s forall a. Num a => a -> a -> a
+ t
1
        in t
s1 seq :: forall a b. a -> b -> b
`seq` case forall a. Integral a => a -> a -> (a, a)
divMod Integer
i Integer
base of
                      (Integer
j,Integer
r)
                        | Integer
r forall a. Eq a => a -> a -> Bool
== Integer
0    -> t -> Integer -> (t, Bool)
exactLoop t
s1 Integer
j
                        | Bool
otherwise -> (forall {t}. Num t => t -> Integer -> t
underLoop t
s1 Integer
j, Bool
False)

  underLoop :: t -> Integer -> t
underLoop t
s Integer
i
    | Integer
i forall a. Ord a => a -> a -> Bool
< Integer
base  = t
s
    | Bool
otherwise = let s1 :: t
s1 = t
s forall a. Num a => a -> a -> a
+ t
1 in t
s1 seq :: forall a b. a -> b -> b
`seq` t -> Integer -> t
underLoop t
s1 (forall a. Integral a => a -> a -> a
div Integer
i Integer
base)


-- | Compute the number of bits required to represent the given integer.
widthInteger :: Integer -> Integer
widthInteger :: Integer -> Integer
widthInteger Integer
x = forall {t} {t}. (Ord t, Bits t, Num t, Num t) => t -> t -> t
go' Integer
0 (if Integer
x forall a. Ord a => a -> a -> Bool
< Integer
0 then forall a. Bits a => a -> a
complement Integer
x else Integer
x)
  where
    go :: t -> t -> t
go t
s t
0 = t
s
    go t
s t
n = let s' :: t
s' = t
s forall a. Num a => a -> a -> a
+ t
1 in t
s' seq :: forall a b. a -> b -> b
`seq` t -> t -> t
go t
s' (t
n forall a. Bits a => a -> Int -> a
`shiftR` Int
1)

    go' :: t -> t -> t
go' t
s t
n
      | t
n forall a. Ord a => a -> a -> Bool
< forall a. Bits a => Int -> a
bit Int
32 = forall {t} {t}. (Num t, Bits t, Num t) => t -> t -> t
go t
s t
n
      | Bool
otherwise  = let s' :: t
s' = t
s forall a. Num a => a -> a -> a
+ t
32 in t
s' seq :: forall a b. a -> b -> b
`seq` t -> t -> t
go' t
s' (t
n forall a. Bits a => a -> Int -> a
`shiftR` Int
32)


-- | Compute the exact root of a natural number.
-- The second argument specifies which root we are computing.
rootExact :: Integer -> Integer -> Maybe Integer
rootExact :: Integer -> Integer -> Maybe Integer
rootExact Integer
x Integer
y = do (Integer
z,Bool
True) <- Integer -> Integer -> Maybe (Integer, Bool)
genRoot Integer
x Integer
y
                   forall (m :: * -> *) a. Monad m => a -> m a
return Integer
z



{- | Compute the the n-th root of a natural number, rounded down to
the closest natural number.  The boolean indicates if the result
is exact (i.e., True means no rounding was done, False means rounded down).
The second argument specifies which root we are computing. -}
genRoot :: Integer -> Integer -> Maybe (Integer, Bool)
genRoot :: Integer -> Integer -> Maybe (Integer, Bool)
genRoot Integer
_  Integer
0    = forall a. Maybe a
Nothing
genRoot Integer
x0 Integer
1    = forall a. a -> Maybe a
Just (Integer
x0, Bool
True)
genRoot Integer
x0 Integer
root = forall a. a -> Maybe a
Just (Integer -> Integer -> (Integer, Bool)
search Integer
0 (Integer
x0forall a. Num a => a -> a -> a
+Integer
1))
  where
  search :: Integer -> Integer -> (Integer, Bool)
search Integer
from Integer
to = let x :: Integer
x = Integer
from forall a. Num a => a -> a -> a
+ forall a. Integral a => a -> a -> a
div (Integer
to forall a. Num a => a -> a -> a
- Integer
from) Integer
2
                       a :: Integer
a = Integer
x forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
root
                   in case forall a. Ord a => a -> a -> Ordering
compare Integer
a Integer
x0 of
                        Ordering
EQ              -> (Integer
x, Bool
True)
                        Ordering
LT | Integer
x forall a. Eq a => a -> a -> Bool
/= Integer
from  -> Integer -> Integer -> (Integer, Bool)
search Integer
x Integer
to
                           | Bool
otherwise  -> (Integer
from, Bool
False)
                        Ordering
GT | Integer
x forall a. Eq a => a -> a -> Bool
/= Integer
to    -> Integer -> Integer -> (Integer, Bool)
search Integer
from Integer
x
                           | Bool
otherwise  -> (Integer
from, Bool
False)