-- |
-- 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 n' :: Nat'
n' =
  case Nat'
n' of
    Nat i :: Integer
i -> Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
i
    _     -> Maybe Integer
forall a. Maybe a
Nothing




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


nAdd :: Nat' -> Nat' -> Nat'
nAdd :: Nat' -> Nat' -> Nat'
nAdd Inf _           = Nat'
Inf
nAdd _ Inf           = Nat'
Inf
nAdd (Nat x :: Integer
x) (Nat y :: 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 0) _       = Integer -> Nat'
Nat 0
nMul _ (Nat 0)       = Integer -> Nat'
Nat 0
nMul Inf _           = Nat'
Inf
nMul _ Inf           = Nat'
Inf
nMul (Nat x :: Integer
x) (Nat y :: 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 0)       = Integer -> Nat'
Nat 1
nExp Inf _           = Nat'
Inf
nExp (Nat 0) Inf     = Integer -> Nat'
Nat 0
nExp (Nat 1) Inf     = Integer -> Nat'
Nat 1
nExp (Nat _) Inf     = Nat'
Inf
nExp (Nat x :: Integer
x) (Nat y :: 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 Inf x :: Nat'
x            = Nat'
x
nMin x :: Nat'
x Inf            = Nat'
x
nMin (Nat x :: Integer
x) (Nat y :: 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 Inf _            = Nat'
Inf
nMax _ Inf            = Nat'
Inf
nMax (Nat x :: Integer
x) (Nat y :: 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 Inf (Nat _)              = Nat' -> Maybe Nat'
forall a. a -> Maybe a
Just Nat'
Inf
nSub (Nat x :: Integer
x) (Nat y :: 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 _ _                      = 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 0)  = Maybe Nat'
forall a. Maybe a
Nothing
nDiv Inf _            = Maybe Nat'
forall a. Maybe a
Nothing
nDiv (Nat x :: Integer
x) (Nat y :: 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 _) Inf      = Nat' -> Maybe Nat'
forall a. a -> Maybe a
Just (Integer -> Nat'
Nat 0)

nMod :: Nat' -> Nat' -> Maybe Nat'
nMod :: Nat' -> Nat' -> Maybe Nat'
nMod _       (Nat 0)  = Maybe Nat'
forall a. Maybe a
Nothing
nMod Inf     _        = Maybe Nat'
forall a. Maybe a
Nothing
nMod (Nat x :: Integer
x) (Nat y :: 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 x :: Integer
x) 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 0)  = Maybe Nat'
forall a. Maybe a
Nothing
nCeilDiv Inf     _        = Maybe Nat'
forall a. Maybe a
Nothing
nCeilDiv (Nat _) Inf      = Maybe Nat'
forall a. Maybe a
Nothing
nCeilDiv (Nat x :: Integer
x) (Nat y :: 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 0)  = Maybe Nat'
forall a. Maybe a
Nothing
nCeilMod Inf     _        = Maybe Nat'
forall a. Maybe a
Nothing
nCeilMod (Nat _) Inf      = Maybe Nat'
forall a. Maybe a
Nothing
nCeilMod (Nat x :: Integer
x) (Nat y :: 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 Inf      = Nat'
Inf
nLg2 (Nat 0)  = Integer -> Nat'
Nat 0
nLg2 (Nat n :: Integer
n)  = case Integer -> Integer -> Maybe (Integer, Bool)
genLog Integer
n 2 of
                  Just (x :: Integer
x,exact :: 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
+ 1)
                  Nothing -> String -> [String] -> Nat'
forall a. HasCallStack => String -> [String] -> a
panic "Cryptol.TypeCheck.Solver.InfNat.nLg2"
                               [ "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 Inf      = Nat'
Inf
nWidth (Nat n :: 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 x :: Integer
x) (Nat y :: Integer
y) (Nat z :: Integer
z)
  | Integer
step Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= 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
+ 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 0 else Integer
len)
                                  -- increasing
                                  else (if Integer
z Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
x then 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 _ _ _ = 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 x :: Integer
x 0    = if Integer
x Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== 1 then (Integer, Bool) -> Maybe (Integer, Bool)
forall a. a -> Maybe a
Just (0, Bool
True) else Maybe (Integer, Bool)
forall a. Maybe a
Nothing
genLog _ 1    = Maybe (Integer, Bool)
forall a. Maybe a
Nothing
genLog 0 _    = Maybe (Integer, Bool)
forall a. Maybe a
Nothing
genLog x :: Integer
x base :: Integer
base = (Integer, Bool) -> Maybe (Integer, Bool)
forall a. a -> Maybe a
Just (Integer -> Integer -> (Integer, Bool)
forall t. Num t => t -> Integer -> (t, Bool)
exactLoop 0 Integer
x)
  where
  exactLoop :: t -> Integer -> (t, Bool)
exactLoop s :: t
s i :: Integer
i
    | Integer
i Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== 1     = (t
s,Bool
True)
    | Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
base   = (t
s,Bool
False)
    | Bool
otherwise  =
        let s1 :: t
s1 = t
s t -> t -> t
forall a. Num a => a -> a -> a
+ 1
        in t
s1 t -> (t, Bool) -> (t, Bool)
forall a b. a -> b -> b
`seq` case Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
divMod Integer
i Integer
base of
                      (j :: Integer
j,r :: Integer
r)
                        | Integer
r Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== 0    -> t -> Integer -> (t, Bool)
exactLoop t
s1 Integer
j
                        | Bool
otherwise -> (t -> Integer -> t
forall t. Num t => t -> Integer -> t
underLoop t
s1 Integer
j, Bool
False)

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

    go' :: t -> t -> t
go' s :: t
s n :: t
n
      | t
n t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> t
forall a. Bits a => Int -> a
bit 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
+ 32 in t
s' t -> t -> t
forall a b. a -> b -> b
`seq` t -> t -> t
go' t
s' (t
n t -> Int -> t
forall a. Bits a => a -> Int -> a
`shiftR` 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 x :: Integer
x y :: Integer
y = do (z :: Integer
z,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 _  0    = Maybe (Integer, Bool)
forall a. Maybe a
Nothing
genRoot x0 :: Integer
x0 1    = (Integer, Bool) -> Maybe (Integer, Bool)
forall a. a -> Maybe a
Just (Integer
x0, Bool
True)
genRoot x0 :: Integer
x0 root :: Integer
root = (Integer, Bool) -> Maybe (Integer, Bool)
forall a. a -> Maybe a
Just (Integer -> Integer -> (Integer, Bool)
search 0 (Integer
x0Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+1))
  where
  search :: Integer -> Integer -> (Integer, Bool)
search from :: Integer
from to :: 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) 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
                        EQ              -> (Integer
x, Bool
True)
                        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)
                        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)