{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE Safe #-}
-- See Note [-Wincomplete-uni-patterns and irrefutable patterns] in Cryptol.TypeCheck.TypePat
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
module Cryptol.TypeCheck.SimpType where

import Control.Applicative((<|>))
import Cryptol.TypeCheck.Type hiding
  (tSub,tMul,tDiv,tMod,tExp,tMin,tCeilDiv,tCeilMod,tLenFromThenTo)
import Cryptol.TypeCheck.TypePat
import Cryptol.TypeCheck.Solver.InfNat
import Control.Monad(msum,guard)


tRebuild' :: Bool -> Type -> Type
tRebuild' withUser = go
  where
  go ty =
    case ty of
      TUser x xs t
        | withUser    -> TUser x xs (go t)
        | otherwise   -> go t
      TVar _          -> ty
      TRec xs         -> TRec (fmap go xs)
      TNominal nt xs  -> TNominal nt (map go xs)
      TCon tc ts      -> tCon tc (map go ts)

tRebuild :: Type -> Type
tRebuild = tRebuild' True

tCon :: TCon -> [Type] -> Type
tCon tc ts =
  case tc of
    TF f ->
      case (f, ts) of
        (TCAdd, [x, y]) -> tAdd x y
        (TCSub, [x, y]) -> tSub x y
        (TCMul, [x, y]) -> tMul x y
        (TCExp, [x, y]) -> tExp x y
        (TCDiv, [x, y]) -> tDiv x y
        (TCMod, [x, y]) -> tMod x y
        (TCMin, [x, y]) -> tMin x y
        (TCMax, [x, y]) -> tMax x y
        (TCWidth, [x]) -> tWidth x
        (TCCeilDiv, [x, y]) -> tCeilDiv x y
        (TCCeilMod, [x, y]) -> tCeilMod x y
        (TCLenFromThenTo, [x, y, z]) -> tLenFromThenTo x y z
        _ -> TCon tc ts
    _ -> TCon tc ts

-- Normal: constants to the left
tAdd :: Type -> Type -> Type
tAdd x y
  | Just t <- tOp TCAdd (total (op2 nAdd)) [x,y] = t
  | tIsInf x            = tInf
  | tIsInf y            = tInf
  | Just n <- tIsNum x  = addK n y
  | Just n <- tIsNum y  = addK n x
  | Just (n,x1) <- isSumK x = addK n (tAdd x1 y)
  | Just (n,y1) <- isSumK y = addK n (tAdd x y1)
  | Just v <- matchMaybe (do (a,b) <- (|-|) y
                             ((guard (x == b) >> return a)
                              -- added to handle this case: 2 ^^ (1 + h) - 1 == 2 ^^ h - 1 + 2 ^^ h
                              <|> (same x a >>= \x2 -> return (tSub x2 b)))
                              ) = v
  | Just v <- matchMaybe (do (a,b) <- (|-|) x
                             ((guard (b == y) >> return a)
                              <|> (same y a >>= \y2 -> return (tSub y2 b)))
                             ) = v

  | Just v <- matchMaybe (factor <|> same x y <|> swapVars) = v

  | otherwise           = tf2 TCAdd x y
  where
  isSumK t = case tNoUser t of
               TCon (TF TCAdd) [ l, r ] ->
                 do n <- tIsNum l
                    return (n, r)
               _ -> Nothing

  addK 0 t = t
  addK n t | Just (m,b) <- isSumK t = tf2 TCAdd (tNum (n + m)) b
           | Just v <- matchMaybe
                     $ do (a,b) <- (|-|) t
                          (do m <- aNat b
                              return $ case compare n m of
                                         GT -> tAdd (tNum (n-m)) a
                                         EQ -> a
                                         LT -> tSub a (tNum (m-n)))
                            <|>
                            (do m <- aNat a
                                return (tSub (tNum (m+n)) b))
                      = v
           -- K + min a b ~> min (K + a) (K + b)
           | Just v <- matchMaybe
                    $ do (a,b) <- aMin t
                         return $ tMin (tAdd (tNum n) a) (tAdd (tNum n) b)
              = v

           | otherwise              = tf2 TCAdd (tNum n) t

  factor = do (a,b1)  <- aMul x
              (a',b2) <- aMul y
              guard (a == a')
              return (tMul a (tAdd b1 b2))

  same x1 y1 = do guard (x1 == y1)
                  return (tMul (tNum (2 :: Int)) x1)

  swapVars = do a <- aTVar x
                b <- aTVar y
                guard (b < a)
                return (tf2 TCAdd y x)

tSub :: Type -> Type -> Type
tSub x y
  | Just t <- tOp TCSub (op2 nSub) [x,y] = t
  | tIsInf y  = tError (tf2 TCSub x y)

  | tIsInf x = x
    {- This assumes that `y` is finite and not error.  The first should
       follow from the typing on `tSub`, which asserts that the second argument
       is finite and less than the first;  the second should have been handled
       by the first equation above, see `tOp`. -}

  | Just 0 <- yNum = x
  | Just k <- yNum
  , TCon (TF TCAdd) [a,b] <- tNoUser x
  , Just n <- tIsNum a = case compare k n of
                           EQ -> b
                           LT -> tf2 TCAdd (tNum (n - k)) b
                           GT -> tSub b (tNum (k - n))

  | Just v <- matchMaybe (do (a,b) <- anAdd x
                             (guard (a == y) >> return b)
                                <|> (guard (b == y) >> return a))
                       = v

    --    x^^(n+h) - x^^h 
    -- ~> (x^^n * x^^h) - x^^h 
    -- ~> ((x^^n - 1) * x^^h
    -- allows subtraction cancelling to occur when
    -- (2^^h + 2^^h) has been rewritten into 2^^(1+h)
  | Just v <- matchMaybe $ 
      do (x_base,x_exp) <- (|^|) x
         (y_base,y_exp) <- (|^|) y
         guard (x_base == y_base)
         x_exp_sum <- anAdd x_exp
         matchSwap x_exp_sum $ \(h,n) -> 
           do guard (h == y_exp)
              let x_to_n = tExp x_base n
              let lhs = tSub x_to_n (tNum (1 :: Int))
              return $ tMul lhs y
       = v

  | Just v <- matchMaybe (do (a,b) <- (|-|) y
                             return (tSub (tAdd x b) a)) = v

  | otherwise = tf2 TCSub x y
  where
  yNum = tIsNum y



-- Normal: constants to the left
tMul :: Type -> Type -> Type
tMul x y
  | Just t <- tOp TCMul (total (op2 nMul)) [x,y] = t
  | Just n <- tIsNum x  = mulK n y
  | Just n <- tIsNum y  = mulK n x
  | Just v <- matchMaybe swapVars = v
  | otherwise = checkExpMul x y

  where
  mulK 0 _ = tNum (0 :: Int)
  mulK 1 t = t
  mulK n t | TCon (TF TCMul) [a,b] <- t'
           , Just a' <- tIsNat' a = case a' of
                                     Inf   -> t
                                     Nat m -> tf2 TCMul (tNum (n * m)) b
           | TCon (TF TCDiv) [a,b] <- t'
           , Just b' <- tIsNum b
           -- XXX: similar for a = b * k?
           , n == b' = tSub a (tMod a b)
           -- c * c ^ x = c ^ (1 + x)
           | TCon (TF TCExp) [a,b] <- t'
           , Just n' <- tIsNum a
           , n == n' = tf2 TCExp a (tAdd (tNum (1::Int)) b)
           -- c^x * c^y = c ^ (y + x)
           | otherwise = tf2 TCMul (tNum n) t
    where t' = tNoUser t

  swapVars = do a <- aTVar x
                b <- aTVar y
                guard (b < a)
                return (tf2 TCMul y x)

  -- Check if (K^a * K^b) => K^(a + b) otherwise default to standard mul
  checkExpMul s t | TCon (TF TCExp) [a,aExp] <- s
                  , Just a' <- tIsNum a
                  , TCon (TF TCExp) [b,bExp] <- t
                  , Just b' <- tIsNum b
                  , (a' >= 2 && a' == b') = tf2 TCExp a (tAdd aExp bExp)
                  | otherwise = tf2 TCMul x y



tDiv :: Type -> Type -> Type
tDiv x y
  | Just t <- tOp TCDiv (op2 nDiv) [x,y] = t
  | tIsInf x = bad
  | Just 0 <- tIsNum y = bad
  | otherwise = tf2 TCDiv x y
    where bad = tError (tf2 TCDiv x y)

tMod :: Type -> Type -> Type
tMod x y
  | Just t <- tOp TCMod (op2 nMod) [x,y] = t
  | tIsInf x = bad
  | Just 0 <- tIsNum y = bad
  | otherwise = tf2 TCMod x y
    where bad = tError (tf2 TCMod x y)

tCeilDiv :: Type -> Type -> Type
tCeilDiv x y
  | Just t <- tOp TCCeilDiv (op2 nCeilDiv) [x,y] = t
  | tIsInf y = bad
  | Just 0 <- tIsNum y = bad
  | otherwise = tf2 TCCeilDiv x y
    where bad = tError (tf2 TCCeilDiv x y)

tCeilMod :: Type -> Type -> Type
tCeilMod x y
  | Just t <- tOp TCCeilMod (op2 nCeilMod) [x,y] = t
  | tIsInf y = bad
  | Just 0 <- tIsNum x = bad
  | otherwise = tf2 TCCeilMod x y
    where bad = tError (tf2 TCCeilMod x y)


tExp :: Type -> Type -> Type
tExp x y
  | Just t <- tOp TCExp (total (op2 nExp)) [x,y] = t
  | Just 0 <- tIsNum y = tNum (1 :: Int)
    -- If `x = (a ^^ b)`, then `(a ^^ b) ^^ y` simplifies to `a ^^ (b * y)`.
    -- This even holds if `a == 0`, given that `0 ^^ 0 == 1` in Cryptol.
  | TCon (TF TCExp) [a,b] <- tNoUser x = tExp a (tMul b y)
  | otherwise = tf2 TCExp x y


-- Normal: constants to the left
tMin :: Type -> Type -> Type
tMin x y
  | Just t <- tOp TCMin (total (op2 nMin)) [x,y] = t
  | Just n <- tIsNat' x = minK n y
  | Just n <- tIsNat' y = minK n x
  | Just n <- matchMaybe (minPlusK x y <|> minPlusK y x) = n
  | Just n <- matchMaybe $ do (k,a) <- isMinK x
                              return $ minK k (tMin a y)
                          <|>
                          do (k,a) <- isMinK y
                             return $ minK k (tMin x a)
    = n

  | Just n <- matchMaybe $ do (k1,a) <- isAddK x
                              (k2,b) <- isAddK y
                              guard (a == b)
                              return $ tAdd (tNum (min k1 k2)) a
    = n

  | x == y              = x
  -- XXX: min (k + t) t -> t
  | otherwise           = tf2 TCMin x y
  where
  isAddK ty = do (a,b) <- anAdd ty
                 k     <- aNat a
                 return (k,b)

  isMinK ty = do (a,b) <- aMin ty
                 k     <- aNat' a
                 return (k,b)

  minPlusK a b = do (k,r) <- isAddK a
                    guard (k >= 1 && b == r)
                    return b


  minK Inf t      = t
  minK (Nat 0) _  = tNum (0 :: Int)
  minK (Nat k) t
    | TCon (TF TCMin) [a,b] <- t'
    , Just n <- tIsNum a   = tf2 TCMin (tNum (min k n)) b

    | otherwise = tf2 TCMin (tNum k) t
    where t' = tNoUser t

-- Normal: constants to the left
tMax :: Type -> Type -> Type
tMax x y
  | Just t <- tOp TCMax (total (op2 nMax)) [x,y] = t
  | Just n <- tIsNat' x = maxK n y
  | Just n <- tIsNat' y = maxK n x
  | otherwise           = tf2 TCMax x y
  where
  maxK Inf _     = tInf
  maxK (Nat 0) t = t
  maxK (Nat k) t

    -- max 1 t ~> t,   if t = a ^ b && a >= 1
    | k == 1
    , TCon (TF TCExp) [a,_] <- t'
    , Just base <- tIsNat' a
    , base >= Nat 1 = t

    | TCon (TF TCAdd) [a,b] <- t'
    , Just n <- tIsNum a = if k <= n
                             then t
                             else tAdd (tNum n) (tMax (tNum (k - n)) b)

    | TCon (TF TCSub) [a,b] <- t'
    , Just n <- tIsNat' a =
      case n of
        Inf   -> t
        Nat m -> if k >= m then tNum k else tSub a (tMin (tNum (m - k)) b)

    | TCon (TF TCMax) [a,b] <- t'
    , Just n <- tIsNum a  = tf2 TCMax (tNum (max k n)) b

    | otherwise = tf2 TCMax (tNum k) t
    where t' = tNoUser t


tWidth :: Type -> Type
tWidth x
  | Just t <- tOp TCWidth (total (op1 nWidth)) [x] = t

  -- width (2^n - 1) = n
  | TCon (TF TCSub) [a,b] <- tNoUser x
  , Just 1 <- tIsNum b
  , TCon (TF TCExp) [p,q] <- tNoUser a
  , Just 2 <- tIsNum p = q

  -- width (2^n) = n + 1
  | TCon (TF TCExp) [m,n] <- tNoUser x
  , Just 2 <- tIsNum m = tf2 TCAdd n (tNum (1 :: Int))

  | otherwise = tf1 TCWidth x

tLenFromThenTo :: Type -> Type -> Type -> Type
tLenFromThenTo x y z
  | Just t <- tOp TCLenFromThenTo (op3 nLenFromThenTo) [x,y,z] = t
  | otherwise = tf3 TCLenFromThenTo x y z

total :: ([Nat'] -> Nat') -> ([Nat'] -> Maybe Nat')
total f xs = Just (f xs)

op1 :: (a -> b) -> [a] -> b
op1 f ~[x] = f x

op2 :: (a -> a -> b) -> [a] -> b
op2 f ~[x,y] = f x y

op3 :: (a -> a -> a -> b) -> [a] -> b
op3 f ~[x,y,z] = f x y z

-- | Common checks: check for error, or simple full evaluation.
-- We assume that input kinds and the result kind are the same (i.e., Nat)
tOp :: TFun -> ([Nat'] -> Maybe Nat') -> [Type] -> Maybe Type
tOp tf f ts
  | Just t <- msum (map tIsError ts) = Just (tError t)
    -- assumes result kind the same as input kind

  | Just xs <- mapM tIsNat' ts =
      Just $ case f xs of
               Nothing -> tError (TCon (TF tf) (map tNat' xs))
               Just n  -> tNat' n
  | otherwise = Nothing




