    -- |
    -- Module      :  $Header$
    -- Copyright   :  (c) 2014-2016 Galois, Inc.
    -- License     :  BSD3
    -- Maintainer  :  cryptol@galois.com
    -- Stability   :  provisional
    -- Portability :  portable
    --
    -- Separate Non-Linear Constraints
    -- When we spot a non-linear expression, we name it and add it to a map.
    --
    -- If we see the same expression multiple times, then we give it the same name.
    --
    -- The body of the non-linear expression is not processed further,
    -- so the resulting map should not contain any of the newly minted names.

    {-# LANGUAGE Safe, RecordWildCards #-}

    module Cryptol.TypeCheck.Solver.Numeric.NonLin
      ( nonLinProp
      , NonLinS, nonLinSubst
      , initialNonLinS
      , apSubstNL
      , lookupNL
      ) where

    import Cryptol.TypeCheck.Solver.Numeric.AST
    import Cryptol.TypeCheck.Solver.Numeric.Simplify
    import Cryptol.TypeCheck.Solver.Numeric.SimplifyExpr
    import Cryptol.Utils.Panic(panic)

    -- import           Data.GenericTrie (Trie)
    -- import qualified Data.GenericTrie as Trie
    import           Data.Maybe ( fromMaybe )
    import           MonadLib
    import           Data.Map (Map)
    import qualified Data.Map as Map

    type Trie   = Map

    trie_empty :: Map k a
    trie_empty  = Map.empty

    trie_insert :: Expr -> a -> Map Expr a -> Map Expr a
    trie_insert = Map.insert

    trie_delete :: Expr -> Map Expr a -> Map Expr a
    trie_delete = Map.delete

    trie_lookup :: Expr -> Map Expr a -> Maybe a
    trie_lookup = Map.lookup

    -- | Factor-out non-linear terms, by naming them.
    nonLinProp :: NonLinS -> Prop -> ([Prop], NonLinS)
    nonLinProp s0 prop = (p : ps, s)
      where ((p,ps),s) = runNL s0 (nonLinPropM prop)

    {- | Apply a substituin to the non-linear expression database.
    Returns `Nothing` if nothing was affected.
    Otherwise returns `Just`, and a substitution for non-linear expressions
    that became linear.

    The definitions of NL terms do not contain other named NL terms,
    so it does not matter if the substitution contains bindings like @_a = e@.

    There should be no bindings that mention NL term names in the definitions
    of the substition (i.e, things like @x = _a@ are NOT ok).
    -}
    apSubstNL :: Subst -> NonLinS -> Maybe (Subst, [Prop], NonLinS)
    apSubstNL su s0 = case runNL s0 (mApSubstNL su) of
                        ((Nothing,_),_) -> Nothing
                        ((Just su1,ps),r) -> Just (su1,ps,r)

    lookupNL :: Name -> NonLinS -> Maybe Expr
    lookupNL x NonLinS { .. } = Map.lookup x nonLinExprs

    runNL :: NonLinS -> NonLinM a -> ((a, [Prop]), NonLinS)
    runNL s m = runId
              $ runStateT s
              $ do a  <- m
                   ps <- finishTodos
                   return (a,ps)

    -- | Get the known non-linear terms.
    nonLinSubst :: NonLinS -> Subst
    nonLinSubst = nonLinExprs

    -- | The initial state for the linearization process.
    initialNonLinS :: NonLinS
    initialNonLinS = NonLinS
      { nextName = 0
      , nonLinExprs = Map.empty
      , nlKnown = trie_empty
      , nlTodo = []
      }

    data SubstOneRes = NoChange
                       -- ^ Substitution does not affect the expression.

                     | Updated (Maybe (Name,Expr))
                       -- ^ The expression was updated and, maybe, it became linear.

    {- | Apply the substituint to all non-linear bindings.
    Returns `Nothing` if nothing was affected.
    Otherwise returns `Just`, and a substituion mapping names that used
    to be non-linear but became linear.

    Note that we may return `Just empty`, indicating that some non-linear
    expressions were updated, but they remained non-linear. -}
    mApSubstNL :: Subst -> NonLinM (Maybe Subst)
    mApSubstNL su =
      do s <- get
         answers <- mapM (mApSubstOneNL su) (Map.toList (nonLinExprs s))
         return (foldr upd Nothing answers)
      where
      upd NoChange ch     = ch
      upd (Updated mb) ch = let lsu = fromMaybe Map.empty ch
                            in Just (case mb of
                                       Nothing    -> lsu
                                       Just (x,e) -> Map.insert x e lsu)

    mApSubstOneNL :: Subst -> (Name,Expr) -> NonLinM SubstOneRes
    mApSubstOneNL su (x,e) =
      case apSubst su e of
        Nothing -> return NoChange
        Just e1 ->
          case crySimpExprMaybe e1 of

            Nothing ->
              sets $ \NonLinS { .. } ->
                ( Updated Nothing
                , NonLinS { nonLinExprs = Map.insert x e1 nonLinExprs
                          , nlKnown = trie_insert e1 x (trie_delete e nlKnown)
                          , .. }
                )

            Just e2
              | isNonLinOp e2 ->
              sets $ \NonLinS { .. } ->
                (Updated Nothing
                , NonLinS { nonLinExprs = Map.insert x e2 nonLinExprs
                          , nlKnown = trie_insert e2 x (trie_delete e nlKnown)
                          , .. }
                )

              | otherwise ->
                do sets_ $ \NonLinS { .. } ->
                     NonLinS { nonLinExprs = Map.delete x nonLinExprs
                             , nlKnown = trie_delete e nlKnown
                             , ..
                             }
                   es <- mapM nonLinExprM (cryExprExprs e2)
                   let e3 = cryRebuildExpr e2 es
                   return (Updated (Just (x,e3)))

    -- | Is the top-level operator a non-linear one.
    isNonLinOp :: Expr -> Bool
    isNonLinOp expr =
      case expr of
        K _   -> False
        Var _ -> False

        _ :+ _ -> False

        _ :- _ -> False

        x :* y ->
          case (x,y) of
            (K _, _)  -> False
            (_, K _)  -> False
            _         -> True

        Div _ y ->
          case y of
            K (Nat n) -> n == 0
            _         -> True

        Mod _ y ->
          case y of
            K (Nat n) -> n == 0
            _         -> True

        _ :^^ _       -> True

        Min _ _       -> False

        Max _ _       -> False

        Width _       -> True

        LenFromThen _ _ _ -> True -- See also comment on `LenFromThenTo`

        LenFromThenTo x y _ ->
          case (x,y) of
            (K _, K _) -> False
            _          -> True    -- Actually, as long as the difference bettwen
                                  -- `x` and `y` is constant we'd be OK, but not
                                  -- sure how to do that...

    nlImplied :: Expr -> Name -> [Prop]
    nlImplied expr x =
      map crySimplify $
      case expr of
        K (Nat n) :^^ e | n >= 2 -> [ Var x :>= one, Var x :>= e :+ one ]
        Mod _ e                  -> [ e     :>= Var x :+ one ]

        _                        -> []

    nonLinPropM :: Prop -> NonLinM Prop
    nonLinPropM prop =
      case prop of
        PFalse      -> return PFalse
        PTrue       -> return PTrue
        Not p       -> Not   `fmap` nonLinPropM p
        p :&& q     -> (:&&) `fmap` nonLinPropM p `ap` nonLinPropM q
        p :|| q     -> (:||) `fmap` nonLinPropM p `ap` nonLinPropM q
        Fin (Var _) -> return prop
        Fin _       -> unexpected
        x :==: y    -> (:==:) `fmap` nonLinExprM x `ap` nonLinExprM y
        x :>: y     -> (:>:)  `fmap` nonLinExprM x `ap` nonLinExprM y

        _ :== _     -> unexpected
        _ :>= _     -> unexpected
        _ :> _      -> unexpected

      where
      unexpected = panic "nonLinPropM" [ show (ppProp prop) ]

    nonLinExprM :: Expr -> NonLinM Expr
    nonLinExprM expr
      | isNonLinOp expr = nameExpr expr
      | otherwise = cryRebuildExpr expr `fmap` mapM nonLinExprM (cryExprExprs expr)

    type NonLinM = StateT NonLinS Id

    data NonLinS = NonLinS
      { nextName    :: !Int
      , nonLinExprs :: Subst
      , nlKnown     :: Trie Expr Name
      , nlTodo      :: [Prop]
      } deriving Show

    nameExpr :: Expr -> NonLinM Expr
    nameExpr e = sets $ \s ->
      case trie_lookup e (nlKnown s) of
        Just x -> (Var x, s)
        Nothing ->
          let x  = nextName s
              n  = SysName x
              s1 = NonLinS { nextName = 1 + x
                           , nonLinExprs = Map.insert n e (nonLinExprs s)
                           , nlKnown = trie_insert e n (nlKnown s)
                           , nlTodo = nlImplied e n ++ nlTodo s
                           }
          in (Var n, s1)

    finishTodos :: NonLinM [Prop]
    finishTodos =
      do s <- get
         case nlTodo s of
           [] -> return []
           p : ps ->
            do set s { nlTodo = ps }
               p'  <- nonLinPropM p
               ps' <- finishTodos
               return (p' : ps')