{-  CAO Compiler
    Copyright (C) 2014 Cryptography and Information Security Group, HASLab - INESC TEC and Universidade do Minho

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>. -}

{- |
Module      :  $Header$
Description :  Sizes of the representation of CAO types.
Copyright   :  (C) 2014 Cryptography and Information Security Group, HASLab - INESC TEC and Universidade do Minho
License     :  GPL

Maintainer  :  Paulo Silva <paufil@di.uminho.pt>
Stability   :  experimental
Portability :  non-portable

This module provides the mechanism to verify the sizes of the representation
for CAO types. The specification file can give static bounds for the type of
the library, allowing the verification of literals before translation.
-}

module Language.CAO.Platform.Literals where

import Data.Bits
import Data.List (genericSplitAt, genericReplicate)

import Language.CAO.Common.Literal
import Language.CAO.Common.Monad
import Language.CAO.Common.Polynomial
import Language.CAO.Common.Var

import Language.CAO.Index
import Language.CAO.Index.Eval

import Language.CAO.Type

{- |
Size of a data type
-}
data Size
    = Generic                 -- ^ Not defined (don't care or unbounded)
    | Simple Integer          -- ^ Size of types other than matrices
    | MSize  Integer Integer  -- ^ Size of a matrix (row, column)
    deriving (Eq, Show)

--------------------------------------------------------------------------------
-- Integer usage
-- Int  : 5
-- Bits : 5
-- Mod  : 4
data LitCheck = LitCheck {
      dynamic :: Bool       -- Bits, Mod (Int is always dynamic)
    , splitC   :: Bool
    , vi1 :: Integer
    , vi2 :: Integer
    , vi3 :: Integer
    , vi4 :: Integer
    , vi5 :: Integer
    {-
    dynamic :: Bool,        -- Bits, Mod (Int is always dynamic)
    min :: Integer,         -- Int (value)
    max :: Integer,         -- Int (value), Bits (size), Mod (module)
    shiftA :: Int,          -- Int, Mod
    mask :: Integer,        -- Int, Mod
    nrChunks :: Int,        -- Int, Mod, Bits
    nrPadd :: Int,          -- Bits
    remBits :: Int,         -- Bits
    chunkSize :: Integer    -- Bits -}
 } deriving Show


emptyLitCheck :: LitCheck
emptyLitCheck = LitCheck False False 0 0 0 0 0

setDynamic :: LitCheck -> LitCheck
setDynamic l = l { dynamic = True }

dynamicCheck :: LitCheck -> Bool
dynamicCheck = dynamic

setSplit :: LitCheck -> LitCheck
setSplit l = l { splitC = True }

splitLit :: LitCheck -> Bool
splitLit = splitC

-- Integers
lowerLimSet :: Integer -> LitCheck -> LitCheck
lowerLimSet n l = l { vi1 = n }

lowerLim :: LitCheck -> Integer
lowerLim = vi1

upperLimSet :: Integer -> LitCheck -> LitCheck
upperLimSet n l = l { vi2 = n }

upperLim :: LitCheck -> Integer
upperLim = vi2

shiftASet :: Integer -> LitCheck -> LitCheck
shiftASet n l = l { vi3 = n }

shiftA :: LitCheck -> Integer
shiftA = vi3

-- Integers and Mods
maskSet :: Integer -> LitCheck -> LitCheck
maskSet n l = l { vi4 = n }

mask :: LitCheck -> Integer
mask = vi4

nrChunksSet :: Integer -> LitCheck -> LitCheck
nrChunksSet n l = l { vi5 = n }

nrChunks :: LitCheck -> Integer
nrChunks = vi5

-- Mods

maxModSet :: Integer -> LitCheck -> LitCheck
maxModSet n l = l { vi2 = n }

maxMod :: LitCheck -> Integer
maxMod = vi2

-- Bit strings

maxSizeSet :: Integer -> LitCheck -> LitCheck
maxSizeSet n l = l { vi1 = n }

maxSize :: LitCheck -> Integer
maxSize = vi1

chunkBitSizeSet :: Integer -> LitCheck -> LitCheck
chunkBitSizeSet n l = l { vi2 = n}

chunkBitSize :: LitCheck -> Integer
chunkBitSize = vi2

nrTypChunkSet :: Integer -> LitCheck -> LitCheck
nrTypChunkSet n l = l { vi3 = n }

nrTypChunk :: LitCheck -> Integer
nrTypChunk = vi3

-- Always disjoint with nrTypChunk
nrChunkSet :: Integer -> LitCheck -> LitCheck
nrChunkSet n l = l { vi3 = n }

nrChunk :: LitCheck -> Integer
nrChunk = vi3

remBitsSet :: Integer -> LitCheck -> LitCheck
remBitsSet n l = l { vi4 = n }

remBits :: LitCheck -> Integer
remBits = vi4

nrPaddChunkSet :: Integer -> LitCheck -> LitCheck
nrPaddChunkSet n l = l { vi5 = n }

nrPaddChunk :: LitCheck -> Integer
nrPaddChunk = vi5

--------------------------------------------------------------------------------
type WordSize      = Int
type NumberOfWords = Int
type WordsPerChunk = Int

onlyPositiveLim :: (Num a, Bits a) => WordSize -> NumberOfWords -> a
-- shift 2 (wordSize * numberOfWords - 1) == 2 ^ (wordSize * numberOfWords)
onlyPositiveLim wordSize numberOfWords = shift 2 (wordSize * numberOfWords - 1)
-- twosComplementLim nrWr = shift 2 (wordSize * nrWr - 2) -- TODO: expand specification

splitNum :: WordSize -> NumberOfWords -> WordsPerChunk -> LitCheck -> LitCheck
splitNum wordSize numberOfWords wordsPerChunk = let
    sa = wordSize * wordsPerChunk
    in   shiftASet (toInteger sa)
       . maskSet (shift 2 (sa - 1) - 1)
       . nrChunksSet (toInteger $ numberOfWords `div` wordsPerChunk)
       . setSplit 

--------------------------------------------------------------------------------
checkInt 
    :: Maybe WordSize -> Maybe (NumberOfWords, Maybe WordsPerChunk) 
    -> Maybe LitCheck
-- Nothing case: Assume that the size is irrelevant or unlimited
checkInt Nothing = const Nothing
checkInt (Just wordSize) = fmap (flip worker emptyLitCheck)
    where
    worker :: (NumberOfWords, Maybe WordsPerChunk) -> LitCheck -> LitCheck
    worker (numberOfWords, wordsPerChunk) = let 
            lim   = onlyPositiveLim wordSize numberOfWords
            other = maybe id (splitNum wordSize numberOfWords) wordsPerChunk
        in other . upperLimSet (lim - 1) . lowerLimSet 0 . setDynamic 

-- XXX: Is there a difference between signed and unsigned bits?
checkBits 
    :: Maybe WordSize -> Size -> Maybe (NumberOfWords, Maybe WordsPerChunk) 
    -> Maybe LitCheck
-- Nothing case: Assume that the size is irrelevant or unlimited
checkBits Nothing _ = const Nothing
checkBits (Just wordSize) strLength = maybe Nothing worker
    where
    worker (numberOfWords, mbWrChunk) = checkStaticSize
        aux
        (Just . maybe id auxChunkGeneric mbWrChunk . maxSizeSet lim)
        lim strLength
        where
        aux litcheck strLen = fmap 
            (\x -> auxChunkSimple x strLen litcheck) mbWrChunk

        auxChunkGeneric wordsPerChunk = 
            nrChunksSet (numberOfChunks wordsPerChunk) . auxChunk wordsPerChunk
        auxChunkSimple wordsPerChunk strLen = let
                len = chunkLength wordsPerChunk
                (nrCompleteChunks, nrPaddChunks, nrPaddBits) = 
                    bsParam (numberOfChunks wordsPerChunk) len strLen
             in   nrTypChunkSet nrCompleteChunks
                . remBitsSet nrPaddBits
                . nrPaddChunkSet nrPaddChunks 
                . auxChunk wordsPerChunk
        auxChunk wordsPerChunk = 
              chunkBitSizeSet (chunkLength wordsPerChunk) . setSplit 

        chunkLength    = toInteger . (wordSize *)
        numberOfChunks = toInteger . div numberOfWords
        lim = toInteger $ wordSize * numberOfWords

bsParam :: Integer -> Integer -> Integer -> (Integer, Integer, Integer)
bsParam numberOfChunks chunkLength strLen = let
        (nrCompleteChunks, nrRemBits) = divMod strLen chunkLength
        nrPaddBits = if nrRemBits == 0 then 0 else chunkLength - nrRemBits
        nrPaddChunks   = numberOfChunks 
                       - nrCompleteChunks
                       - (if nrPaddBits == 0 then 0 else 1)
    in (nrCompleteChunks, nrPaddChunks, nrPaddBits)

checkMod 
    :: Maybe WordSize -> Size -> Maybe (NumberOfWords, Maybe WordsPerChunk) 
    -> Maybe LitCheck
checkMod Nothing _ = const Nothing
-- Nothing case: Assume that the size is irrelevant or unlimited
checkMod (Just wordSize) strLength = maybe Nothing worker
    where
    worker :: (NumberOfWords, Maybe WordsPerChunk) -> Maybe LitCheck
    worker (numberOfWords, mbWrChunk) = let
            auxChunk = splitNum wordSize numberOfWords
            lim      = onlyPositiveLim wordSize numberOfWords
        in checkStaticSize
            (\litcheck -> const $ fmap (flip auxChunk litcheck) mbWrChunk)
            (Just . maybe id auxChunk mbWrChunk . maxModSet lim)
            lim strLength

{-
When the size of the type is defined statically and we have the maximum limit,
we can determine if the representation is enough to hold all the values of the
type. Otherwise, the compilation fails.
If the size is not defined, then dinamic verification is set. In this context,
dynamic does not mean during execution, but during translation, i.e., all 
literal values are verified for the bounds. In the future, this can also be
used in an abstract analysis of all values.
-}
checkStaticSize 
    :: (LitCheck -> Integer -> Maybe LitCheck) 
    -> (LitCheck -> Maybe LitCheck) 
    -> Integer -> Size 
    -> Maybe LitCheck
checkStaticSize fsimple fgeneric lim strLength =
    case strLength of
        -- The size can be determined statically
        Simple strLen -> if strLen <= lim
            then fsimple emptyLitCheck strLen 
            else error $ "The platform type `" -- ++ nameInPlat ts ++ "' does not have \
                   -- \ enough precision to hold mods with module " ++ show strLength
        -- No static information about the bit strings.
        -- Dynamic testing required
        Generic -> fgeneric $ setDynamic emptyLitCheck
        _ -> error "Not expected matrix size"

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

checkILit 
    :: CaoMonad m 
    => LitCheck -> Integer 
    -> m [Integer]
checkILit lspec num = 
    if lowerLim lspec <= num && num <= upperLim lspec
        then if splitLit lspec
            then return $ chopNumLit (nrChunks lspec) (mask lspec) (fromInteger $ shiftA lspec) num
            else return [num]
        else error $ "Literal outside of the range of representation: " ++ show num

chopNumLit :: Integer -> Integer -> Int -> Integer -> [Integer]
chopNumLit 0 _ _ _ = []
chopNumLit n msk shft num = let
        cn = num .&. msk
        num' = num `shiftR` shft
    in chopNumLit (n-1) msk shft num' ++ [cn] -- Endianness

checkBSLit :: CaoMonad m => LitCheck -> Type Var -> [Bool] -> m [[Bool]]
checkBSLit lspec (Bits sign ilen) bits =
    case evalExpr ilen of
        IInt len -> if dynamicCheck lspec
            then if len <= maxSize lspec
                then if splitLit lspec
                    then let chunkLength = chunkBitSize lspec
                             (nrCompleteChunks, nrPaddBits, nrPaddChunks) = bsParam (nrChunks lspec) chunkLength len
                        in return $ chopBSLit sign chunkLength nrCompleteChunks nrPaddChunks nrPaddBits bits
                    else return [bits]
                else error $ "Literal outside of the range of representation: " -- ++ bits
            -- If it is not dynamic AND the specification exists, it MUST be choked
            else return $ chopBSLit sign (chunkBitSize lspec) (nrTypChunk lspec) (nrPaddChunk lspec) (nrPaddChunk lspec) bits
        _ -> error $ "Not expected index"
checkBSLit _ _ _ = error $ "Not expected type"

chopBSLit :: Sign -> Integer -> Integer -> Integer -> Integer -> [Bool] -> [[Bool]]
chopBSLit sign chunkLength nrCompleteChunks nrPaddChunks nrPaddBits bits =
    worker [] bits nrCompleteChunks 
    where
    worker acc rest 0 = padd acc rest
    worker acc bstr nc = let
            (pref, rest) = genericSplitAt chunkLength bstr
        in worker (pref : acc) rest (nc - 1)

    padd acc bstr = let
        ebit = case sign of
                S -> last 
                U -> const False
        in if nrPaddBits == 0
            then genericReplicate nrPaddChunks (blankChunk (ebit $ head acc)) ++ acc
            else let
                    ebit' = ebit bstr
                in genericReplicate nrPaddChunks (blankChunk ebit') ++
                   (bstr ++ genericReplicate nrPaddBits ebit') :
                   acc
    blankChunk = genericReplicate chunkLength

        

checkPLit 
    :: CaoMonad m 
    => LitCheck -> Type Var -> Pol Var 
    -> m [Pol Var]
checkPLit lspec (Mod Nothing Nothing (Pol [Mon (CoefI idx) EZero]))
                (Pol [Mon (CoefI cexp) EZero]) =
    case (evalExpr idx, evalExpr cexp) of
        (IInt i, IInt c) -> if dynamicCheck lspec
            then if i <= maxMod lspec
                then return (aux c)
                else error $ "Literal outside of the range of representation"
            else return (aux c)
        _ -> error "checkPLit: Not literal!"
    where
    aux = map (\ c' -> Pol [Mon (CoefI (IInt c')) EZero]) . 
        chopNumLit (nrChunks lspec) (mask lspec) (fromInteger $ shiftA lspec)
checkPLit _ _ _ = error $ "Not supported literal of polynomial"