-- |
-- 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
(Int -> Nat' -> ShowS)
-> (Nat' -> String) -> ([Nat'] -> ShowS) -> Show Nat'
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
(Nat' -> Nat' -> Bool) -> (Nat' -> Nat' -> Bool) -> Eq Nat'
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'
Eq Nat'
-> (Nat' -> Nat' -> Ordering)
-> (Nat' -> Nat' -> Bool)
-> (Nat' -> Nat' -> Bool)
-> (Nat' -> Nat' -> Bool)
-> (Nat' -> Nat' -> Bool)
-> (Nat' -> Nat' -> Nat')
-> (Nat' -> Nat' -> Nat')
-> Ord 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
$cp1Ord :: Eq Nat'
Ord, (forall x. Nat' -> Rep Nat' x)
-> (forall x. Rep Nat' x -> Nat') -> Generic Nat'
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' -> ()
(Nat' -> ()) -> NFData 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 -> Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
i
    Nat'
_     -> Maybe Integer
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 Integer -> Integer -> Integer
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 Integer -> Integer -> Integer
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 Integer -> Integer -> Integer
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 (Integer -> Integer -> Integer
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 (Integer -> Integer -> Integer
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
_)              = Nat' -> Maybe Nat'
forall a. a -> Maybe a
Just Nat'
Inf
nSub (Nat Integer
x) (Nat Integer
y)
  | Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
y                    = Nat' -> Maybe Nat'
forall a. a -> Maybe a
Just (Integer -> Nat'
Nat (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
y))
nSub Nat'
_ Nat'
_                      = Maybe 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)  = Maybe Nat'
forall a. Maybe a
Nothing
nDiv Nat'
Inf Nat'
_            = Maybe Nat'
forall a. Maybe a
Nothing
nDiv (Nat Integer
x) (Nat Integer
y)  = Nat' -> Maybe Nat'
forall a. a -> Maybe a
Just (Integer -> Nat'
Nat (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div Integer
x Integer
y))
nDiv (Nat Integer
_) Nat'
Inf      = Nat' -> Maybe Nat'
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)  = Maybe Nat'
forall a. Maybe a
Nothing
nMod Nat'
Inf     Nat'
_        = Maybe Nat'
forall a. Maybe a
Nothing
nMod (Nat Integer
x) (Nat Integer
y)  = Nat' -> Maybe Nat'
forall a. a -> Maybe a
Just (Integer -> Nat'
Nat (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
mod Integer
x Integer
y))
nMod (Nat Integer
x) Nat'
Inf      = Nat' -> Maybe Nat'
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@.
-- It is also undefined when either input is infinite; perhaps this
-- could be relaxed later.
nCeilDiv :: Nat' -> Nat' -> Maybe Nat'
nCeilDiv :: Nat' -> Nat' -> Maybe Nat'
nCeilDiv Nat'
_       (Nat Integer
0)  = Maybe Nat'
forall a. Maybe a
Nothing
nCeilDiv Nat'
Inf     Nat'
_        = Maybe Nat'
forall a. Maybe a
Nothing
nCeilDiv (Nat Integer
_) Nat'
Inf      = Maybe Nat'
forall a. Maybe a
Nothing
nCeilDiv (Nat Integer
x) (Nat Integer
y)  = Nat' -> Maybe Nat'
forall a. a -> Maybe a
Just (Integer -> Nat'
Nat (- Integer -> Integer -> Integer
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@.
-- It is also undefined when either input is infinite; perhaps this
-- could be relaxed later.
nCeilMod :: Nat' -> Nat' -> Maybe Nat'
nCeilMod :: Nat' -> Nat' -> Maybe Nat'
nCeilMod Nat'
_       (Nat Integer
0)  = Maybe Nat'
forall a. Maybe a
Nothing
nCeilMod Nat'
Inf     Nat'
_        = Maybe Nat'
forall a. Maybe a
Nothing
nCeilMod (Nat Integer
_) Nat'
Inf      = Maybe Nat'
forall a. Maybe a
Nothing
nCeilMod (Nat Integer
x) (Nat Integer
y)  = Nat' -> Maybe Nat'
forall a. a -> Maybe a
Just (Integer -> Nat'
Nat (Integer -> Integer -> Integer
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 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)
                  Maybe (Integer, Bool)
Nothing -> String -> [String] -> Nat'
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 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
0 = let len :: Integer
len = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div Integer
dist Integer
step Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1
                in Nat' -> Maybe Nat'
forall a. a -> Maybe a
Just (Nat' -> Maybe Nat') -> Nat' -> Maybe Nat'
forall a b. (a -> b) -> a -> b
$ Integer -> Nat'
Nat (Integer -> Nat') -> Integer -> Nat'
forall a b. (a -> b) -> a -> b
$ if Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
y
                                  -- decreasing
                                  then (if Integer
z Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
x then Integer
0 else Integer
len)
                                  -- increasing
                                  else (if Integer
z Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
x then Integer
0 else Integer
len)
  where
  step :: Integer
step = Integer -> Integer
forall a. Num a => a -> a
abs (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
y)
  dist :: Integer
dist = Integer -> Integer
forall a. Num a => a -> a
abs (Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
z)

nLenFromThenTo Nat'
_ Nat'
_ Nat'
_ = Maybe 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 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1 then (Integer, Bool) -> Maybe (Integer, Bool)
forall a. a -> Maybe a
Just (Integer
0, Bool
True) else Maybe (Integer, Bool)
forall a. Maybe a
Nothing
genLog Integer
_ Integer
1    = Maybe (Integer, Bool)
forall a. Maybe a
Nothing
genLog Integer
0 Integer
_    = Maybe (Integer, Bool)
forall a. Maybe a
Nothing
genLog Integer
x Integer
base = (Integer, Bool) -> Maybe (Integer, Bool)
forall a. a -> Maybe a
Just (Integer -> Integer -> (Integer, Bool)
forall a. Num a => a -> Integer -> (a, Bool)
exactLoop Integer
0 Integer
x)
  where
  exactLoop :: a -> Integer -> (a, Bool)
exactLoop a
s Integer
i
    | Integer
i Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
1     = (a
s,Bool
True)
    | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
base   = (a
s,Bool
False)
    | Bool
otherwise  =
        let s1 :: a
s1 = a
s a -> a -> a
forall a. Num a => a -> a -> a
+ a
1
        in a
s1 a -> (a, Bool) -> (a, Bool)
`seq` case Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
divMod Integer
i Integer
base of
                      (Integer
j,Integer
r)
                        | Integer
r Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0    -> a -> Integer -> (a, Bool)
exactLoop a
s1 Integer
j
                        | Bool
otherwise -> (a -> Integer -> a
forall t. Num t => t -> Integer -> t
underLoop a
s1 Integer
j, Bool
False)

  underLoop :: t -> Integer -> t
underLoop t
s Integer
i
    | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
base  = t
s
    | Bool
otherwise = let s1 :: t
s1 = t
s t -> t -> t
forall a. Num a => a -> a -> a
+ t
1 in t
s1 t -> t -> t
`seq` t -> Integer -> t
underLoop t
s1 (Integer -> Integer -> Integer
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 = Integer -> Integer -> Integer
forall t t. (Ord t, Bits t, Num t, Num t) => t -> t -> t
go' Integer
0 (if Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 then Integer -> Integer
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 t -> t -> t
forall a. Num a => a -> a -> a
+ t
1 in t
s' t -> t -> t
`seq` t -> t -> t
go t
s' (t
n t -> Int -> t
forall a. Bits a => a -> Int -> a
`shiftR` Int
1)

    go' :: t -> t -> t
go' t
s t
n
      | t
n t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> t
forall a. Bits a => Int -> a
bit Int
32 = t -> t -> t
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 t -> t -> t
forall a. Num a => a -> a -> a
+ t
32 in t
s' t -> t -> t
`seq` t -> t -> t
go' t
s' (t
n t -> Int -> t
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
                   Integer -> Maybe Integer
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    = Maybe (Integer, Bool)
forall a. Maybe a
Nothing
genRoot Integer
x0 Integer
1    = (Integer, Bool) -> Maybe (Integer, Bool)
forall a. a -> Maybe a
Just (Integer
x0, Bool
True)
genRoot Integer
x0 Integer
root = (Integer, Bool) -> Maybe (Integer, Bool)
forall a. a -> Maybe a
Just (Integer -> Integer -> (Integer, Bool)
search Integer
0 (Integer
x0Integer -> Integer -> Integer
forall 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 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div (Integer
to Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
from) Integer
2
                       a :: Integer
a = Integer
x Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
root
                   in case Integer -> Integer -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Integer
a Integer
x0 of
                        Ordering
EQ              -> (Integer
x, Bool
True)
                        Ordering
LT | Integer
x Integer -> Integer -> Bool
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 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
to    -> Integer -> Integer -> (Integer, Bool)
search Integer
from Integer
x
                           | Bool
otherwise  -> (Integer
from, Bool
False)