-- |
-- Module      :  $Header$
-- Copyright   :  (c) 2014-2016 Galois, Inc.
-- License     :  BSD3
-- Maintainer  :  cryptol@galois.com
-- Stability   :  provisional
-- Portability :  portable
--
-- The sytnax of numeric propositions.

{-# LANGUAGE Safe #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE DeriveGeneric #-}
module Cryptol.TypeCheck.Solver.Numeric.AST
  ( Name(..), ppName

  , Prop(..), cryPropExprs, cryPropFVS
  , ppProp, ppPropPrec

  , Expr(..), zero, one, two, inf, cryAnds, cryOrs
  , cryExprExprs, cryRebuildExpr
  , cryExprFVS
  , ppExpr, ppExprPrec

  , Nat'(..)

  , IfExpr, IfExpr'(..), ppIf, ppIfExpr

  , Subst, HasVars(..), cryLet, composeSubst, doAppSubst
  ) where

import          Cryptol.TypeCheck.AST(TVar)
import          Cryptol.TypeCheck.PP(pp)
import          Cryptol.TypeCheck.Solver.InfNat ( Nat'(..) )
import          Cryptol.Utils.PP ( Doc, text, (<+>), hang, ($$), char, (<>)
                                 , parens, integer, sep )
import          Cryptol.Utils.Panic ( panic )
import          Cryptol.Utils.Misc ( anyJust )

-- import           Data.GenericTrie (TrieKey)
import           GHC.Generics(Generic)
import           Data.Maybe (fromMaybe)
import           Data.Map ( Map )
import qualified Data.Map as Map
import           Data.Set ( Set )
import qualified Data.Set as Set
import qualified Control.Applicative as A
import           Control.Monad ( liftM, ap )


infixr 2 :||
infixr 3 :&&
infix  4 :==, :>, :>=, :==:, :>:
infixl 6 :+, :-
infixl 7 :*
infixr 8 :^^



data Name = UserName TVar | SysName Int
            deriving (Show,Eq,Ord,Generic)



-- | Propopsitions, representing Cryptol's numeric constraints (and a bit more).
data Prop =

   -- Preidcates on natural numbers with infinity.
   -- After simplification, the only one of these should be `fin x`,
   -- where `x` is a variable.

   Fin Expr | Expr :== Expr | Expr :>= Expr | Expr :> Expr


  -- Predicate on strict natural numbers (i.e., no infinities)
  -- Should be introduced by 'cryNatOp', to eliminte 'inf'.
  | Expr :==: Expr | Expr :>: Expr

  -- Standard logical strucutre>
  | Prop :&& Prop | Prop :|| Prop
  | Not Prop
  | PFalse | PTrue
    deriving (Eq,Show,Generic)


-- | Expressions, representing Cryptol's numeric types.
data Expr = K Nat'
          | Var Name
          | Expr :+ Expr                  -- total
          | Expr :- Expr                  -- partial: x >= y, fin y
          | Expr :* Expr                  -- total
          | Div Expr Expr                 -- partial: fin x, y >= 1
          | Mod Expr Expr                 -- partial: fin x, y >= 1
          | Expr :^^ Expr                 -- total
          | Min Expr Expr                 -- total
          | Max Expr Expr                 -- total
          | Width Expr                    -- total
          | LenFromThen   Expr Expr Expr  -- partial: x /= y, w >= width x
          | LenFromThenTo Expr Expr Expr  -- partial: x /= y
            deriving (Eq,Show,Generic,Ord)


-- | The constant @0@.
zero :: Expr
zero = K (Nat 0)

-- | The constant @1@.
one :: Expr
one = K (Nat 1)

-- | The constant @2@.
two :: Expr
two = K (Nat 2)

-- | The constant @infinity@.
inf :: Expr
inf = K Inf


-- | Make a conjucntion of the given properties.
cryAnds :: [Prop] -> Prop
cryAnds []  = PTrue
cryAnds ps  = foldr1 (:&&) ps

-- | Make a disjunction of the given properties.
cryOrs :: [Prop] -> Prop
cryOrs []   = PFalse
cryOrs ps   = foldr1 (:||) ps




-- | Compute all expressions in a property.
cryPropExprs :: Prop -> [Expr]
cryPropExprs = go []
  where
  go es prop =
    case prop of
      PTrue     -> es
      PFalse    -> es
      Not p     -> go es p
      p :&& q   -> go (go es q) p
      p :|| q   -> go (go es q) p

      Fin x     -> x : es

      x :== y   -> x : y : es
      x :>  y   -> x : y : es
      x :>= y   -> x : y : es

      x :==: y  -> x : y : es
      x :>:  y  -> x : y : es


-- | Compute the immediate sub-expressions of an expression.
cryExprExprs :: Expr -> [Expr]
cryExprExprs expr =
  case expr of
    K _                 -> []
    Var _               -> []
    x :+ y              -> [x,y]
    x :- y              -> [x,y]
    x :* y              -> [x,y]
    Div x y             -> [x,y]
    Mod x y             -> [x,y]
    x :^^ y             -> [x,y]
    Min x y             -> [x,y]
    Max x y             -> [x,y]
    Width x             -> [x]
    LenFromThen   x y z -> [x,y,z]
    LenFromThenTo x y z -> [x,y,z]

-- | Rebuild an expression, using the top-level strucutre of the first
-- expression, but the second list of expressions as sub-expressions.
cryRebuildExpr :: Expr -> [Expr] -> Expr
cryRebuildExpr expr args =
  case (expr,args) of
    (K _,   [])                     -> expr
    (Var _, [])                     -> expr
    (_ :+ _k, [x,y])                -> x :+ y
    (_ :- _ , [x,y])                -> x :- y
    (_ :* _ , [x,y])                -> x :* y
    (Div _ _, [x,y])                -> Div x y
    (Mod _ _, [x,y])                -> Mod x y
    (_ :^^ _, [x,y])                -> x :^^ y
    (Min _ _, [x,y])                -> Min x y
    (Max _ _, [x,y])                -> Max x y
    (Width _, [x])                  -> Width x
    (LenFromThen   _ _ _ , [x,y,z]) -> LenFromThen x y z
    (LenFromThenTo _ _ _ , [x,y,z]) -> LenFromThenTo x y z
    _ -> panic "cryRebuildExpr" $ map show
           $ text "expr:" <+> ppExpr expr
           : [ text "arg:" <+> ppExpr a | a <- args ]


-- | Compute the free variables in an expression.
cryExprFVS :: Expr -> Set Name
cryExprFVS expr =
  case expr of
    Var x -> Set.singleton x
    _     -> Set.unions (map cryExprFVS (cryExprExprs expr))

-- | Compute the free variables in a proposition.
cryPropFVS :: Prop -> Set Name
cryPropFVS = Set.unions . map cryExprFVS . cryPropExprs





data IfExpr' p a = If p (IfExpr' p a) (IfExpr' p a) | Return a | Impossible
                deriving Eq

type IfExpr = IfExpr' Prop

instance Monad (IfExpr' p) where
  return  = Return
  fail _  = Impossible
  m >>= k = case m of
              Impossible -> Impossible
              Return a   -> k a
              If p t e   -> If p (t >>= k) (e >>= k)

instance Functor (IfExpr' p) where
  fmap  = liftM

instance A.Applicative (IfExpr' p) where
  pure  = return
  (<*>) = ap


--------------------------------------------------------------------------------
-- Substitution
--------------------------------------------------------------------------------

type Subst = Map Name Expr

composeSubst :: Subst -> Subst -> Subst
composeSubst g f = Map.union f' g
  where
  f' = fmap (\e -> fromMaybe e (apSubst g e)) f

cryLet :: HasVars e => Name -> Expr -> e -> Maybe e
cryLet x e = apSubst (Map.singleton x e)

doAppSubst :: HasVars a => Subst -> a -> a
doAppSubst su a = fromMaybe a (apSubst su a)

-- | Replaces occurances of the name with the expression.
-- Returns 'Nothing' if there were no occurances of the name.
class HasVars ast where
  apSubst :: Subst -> ast -> Maybe ast

-- | This is used in the simplification to "apply" substitutions to Props.
instance HasVars Bool where
  apSubst _ _ = Nothing

instance HasVars Expr where
  apSubst su = go
    where
    go expr =
      case expr of
        K _                 -> Nothing
        Var b               -> Map.lookup b su
        x :+ y              -> bin (:+) x y
        x :- y              -> bin (:-) x y
        x :* y              -> bin (:*) x y
        x :^^ y             -> bin (:^^) x y
        Div x y             -> bin Div x y
        Mod x y             -> bin Mod x y
        Min x y             -> bin Min x y
        Max x y             -> bin Max x y
        Width x             -> Width `fmap` go x
        LenFromThen x y w   -> three LenFromThen x y w
        LenFromThenTo x y z -> three LenFromThen x y z

    bin f x y = do [x',y'] <- anyJust go [x,y]
                   return (f x' y')

    three f x y z = do [x',y',z'] <- anyJust go [x,y,z]
                       return (f x' y' z')

instance HasVars Prop where
  apSubst su = go
    where
    go prop =
      case prop of
        PFalse    -> Nothing
        PTrue     -> Nothing
        Not p     -> Not `fmap` go p
        p :&& q   -> bin (:&&) p q
        p :|| q   -> bin (:||) p q
        Fin x     -> Fin `fmap` apSubst su x
        x :== y   -> twoE (:==) x y
        x :>= y   -> twoE (:>=) x y
        x :> y    -> twoE (:>) x y
        x :==: y  -> twoE (:==:) x y
        x :>: y   -> twoE (:>) x y

    bin f x y = do [x',y'] <- anyJust go [x,y]
                   return (f x' y')

    twoE f x y = do [x',y'] <- anyJust (apSubst su) [x,y]
                    return (f x' y')


--------------------------------------------------------------------------------
-- Tries
--------------------------------------------------------------------------------

-- instance TrieKey Name
-- instance TrieKey Prop
-- instance TrieKey Expr




--------------------------------------------------------------------------------
-- Pretty Printing
--------------------------------------------------------------------------------

-- | Pretty print a name.
ppName :: Name -> Doc
ppName name =
  case name of
    UserName x -> pp x
    SysName  x -> char '_' <> text (names !! x)

-- | An infinite list of names, for pretty prinitng.
names :: [String]
names  = concatMap gen [ 0 :: Integer .. ]
  where
  gen x  = [ a : suff x | a <- [ 'a' .. 'z' ] ]

  suff 0 = ""
  suff x = show x



-- | Pretty print a top-level property.
ppProp :: Prop -> Doc
ppProp = ppPropPrec 0

-- | Pretty print a proposition, in the given precedence context.
ppPropPrec :: Int -> Prop -> Doc
ppPropPrec prec prop =
  case prop of
    Fin x     -> fun "fin" ppExprPrec x
    x :== y   -> bin "==" 4 1 1 ppExprPrec x y
    x :>= y   -> bin ">=" 4 1 1 ppExprPrec x y
    x :> y    -> bin ">"  4 1 1 ppExprPrec x y

    x :==: y  -> bin "==#" 4 1 1 ppExprPrec x y
    x :>: y   -> bin ">#"  4 1 1 ppExprPrec x y

    p :&& q   -> bin "&&" 3 1 0 ppPropPrec p q
    p :|| q   -> bin "||" 2 1 0 ppPropPrec p q
    Not p     -> fun "not" ppPropPrec p
    PTrue     -> text "True"
    PFalse    -> text "False"

  where
  wrap p d = if prec > p then parens d else d

  fun f how x = wrap 10 (text f <+> how 11 x)

  bin op opP lMod rMod how x y =
    wrap opP (sep [ how (opP + lMod) x, text op, how (opP + rMod) y ])



-- | Pretty print an expression at the top level.
ppExpr :: Expr -> Doc
ppExpr = ppExprPrec 0

-- | Pretty print an expression, in the given precedence context.
ppExprPrec :: Int -> Expr -> Doc
ppExprPrec prec expr =
  case expr of
    K Inf               -> text "inf"
    K (Nat n)           -> integer n
    Var a               -> ppName a
    x :+ y              -> bin "+" 6 0 1 x y
    x :- y              -> bin "-" 6 0 1 x y
    x :* y              -> bin "*" 7 0 1 x y
    Div x y             -> fun "div" [x,y]
    Mod x y             -> fun "mod" [x,y]
    x :^^ y             -> bin "^^" 8 1 0 x y
    Min x y             -> fun "min" [x,y]
    Max x y             -> fun "max" [x,y]
    Width x             -> fun "width" [x]
    LenFromThen x y w   -> fun "lenFromThen" [x,y,w]
    LenFromThenTo x y z -> fun "lenFromThenTo" [x,y,z]

  where
  wrap p d = if prec > p then parens d else d

  fun f xs = wrap 10 (text f <+> sep (map (ppExprPrec 11) xs))

  bin op opP lMod rMod x y =
    wrap opP
      (ppExprPrec (opP + lMod) x <+> text op <+> ppExprPrec (opP + rMod) y)



-- | Pretty print an experssion with ifs.
ppIfExpr :: IfExpr Expr -> Doc
ppIfExpr = ppIf ppProp ppExpr

-- | Pretty print an experssion with ifs.
ppIf :: (p -> Doc) -> (a -> Doc) -> IfExpr' p a -> Doc
ppIf ppAtom ppVal = go
  where
  go expr =
    case expr of
      If p t e -> hang (text "if" <+> ppAtom p) 2
                ( (text "then" <+> go t)  $$
                  (text "else" <+> go e)
                )
      Return e    -> ppVal e
      Impossible  -> text "<impossible>"