content: {-
     Copyright 2016, Dominic Orchard, Andrew Rice, Mistral Contrastin, Matthew Danish
  
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
     You may obtain a copy of the License at
  
         http://www.apache.org/licenses/LICENSE-2.0
  
     Unless required by applicable law or agreed to in writing, software
     distributed under the License is distributed on an "AS IS" BASIS,
     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     See the License for the specific language governing permissions and
     limitations under the License.
  -}
  
  {-
  
  This files gives an executable implementation of the model for
  abstract stencil specifications. This model is used to drive both
  the specification checking and program synthesis features.
  
  -}
  
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE DeriveFunctor #-}
  {-# LANGUAGE DeriveFoldable #-}
  {-# LANGUAGE DeriveTraversable #-}
  {-# LANGUAGE DeriveDataTypeable #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE MultiWayIf #-}
  
  module Camfort.Specification.Stencils.Model ( Interval(..)
                                              , Bound(..)
                                              , approxVec
                                              , Offsets(..)
                                              , UnionNF
                                              , vecLength
                                              , unfCompare
                                              , optimise
                                              , maximas
                                              , Approximation(..)
                                              , lowerBound, upperBound
                                              , fromExact
                                              , Multiplicity(..)
                                              , Peelable(..)
                                              ) where
  
  import           Prelude hiding (pred)
  
  import qualified Control.Monad as CM
  
  import           Algebra.Lattice
  import qualified Data.List.NonEmpty as NE
  import qualified Data.Set as S
  import           Data.Foldable
  import           Data.SBV
  import           Data.Data
  import           Data.List (sortBy, nub)
  import           Data.Maybe (fromJust)
  import qualified Camfort.Specification.Stencils.PartialOrd as PO
  
  import qualified Camfort.Helpers.Vec as V
  import System.IO.Unsafe
  
  -- Utility container
  class Container a where
    type MemberTyp a
    type CompTyp a
  
    member :: MemberTyp a -> a -> Bool
    compile :: a -> (CompTyp a -> SBool)
  
  --------------------------------------------------------------------------------
  -- Arbitrary sets representing offsets
  --------------------------------------------------------------------------------
  
  data Offsets =
      Offsets (S.Set Int64)
    | SetOfIntegers
    deriving Eq
  
  instance Ord Offsets where
      Offsets s `compare` Offsets s' = s `compare` s'
      Offsets _ `compare` SetOfIntegers = LT
      SetOfIntegers `compare` Offsets _ = GT
      SetOfIntegers `compare` SetOfIntegers = EQ
  
  instance Container Offsets where
    type MemberTyp Offsets = Int64
    type CompTyp Offsets = SInt64
  
    member i (Offsets s) = i `S.member` s
    member _ _ = True
  
    compile (Offsets s) i = i `sElem` map fromIntegral (S.toList s)
    compile SetOfIntegers _ = sTrue
  
  instance Lattice Offsets where
    (Offsets s) \/ (Offsets s') = Offsets $ s `S.union` s'
    _ \/ _ = SetOfIntegers
    (Offsets s) /\ (Offsets s') = Offsets $ s `S.intersection` s'
    off@Offsets{} /\ _ = off
    _ /\ o = o
  
  instance BoundedJoinSemiLattice Offsets where
    bottom = Offsets S.empty
  
  instance BoundedMeetSemiLattice Offsets where
    top = SetOfIntegers
  
  --------------------------------------------------------------------------------
  -- Interval as defined in the paper
  --------------------------------------------------------------------------------
  
  data Bound = Arbitrary | Standard
  
  -- | Interval data structure assumes the following:
  -- 1. The first num. param. is less than the second;
  -- 2. For holed intervals, first num. param. <= 0 <= second num. param.;
  data Interval a where
    IntervArbitrary :: Int -> Int -> Interval 'Arbitrary
    IntervInfiniteArbitrary :: Interval 'Arbitrary
    IntervHoled     :: Int64 -> Int64 -> Bool -> Interval 'Standard
    IntervInfinite  :: Interval 'Standard
  
  deriving instance Eq (Interval a)
  
  instance Show (Interval 'Standard) where
    show IntervInfinite = "IntervInfinite"
    show (IntervHoled lb up p) =
      "Interv [" ++ show lb ++ "," ++ show up ++ "]^" ++ show p
  
  approxInterv :: Interval 'Arbitrary -> Approximation (Interval 'Standard)
  approxInterv (IntervArbitrary a b)
    | a > b = error
      "Interval condition violated: lower bound is bigger than the upper bound."
    | a <=  0, b >=  0 = Exact  $ IntervHoled a' b' True
    | a <= -1, b == -1 = Exact  $ IntervHoled a' 0  False
    | a ==  1, b >=  1 = Exact  $ IntervHoled 0  b' False
    | a >   1, b >   1 = Bound Nothing $ Just $ IntervHoled 0  b' False
    | a <  -1, b <  -1 = Bound Nothing $ Just $ IntervHoled a' 0  False
    | otherwise = error "Impossible: All posibilities are covered."
    where
      a' = fromIntegral a
      b' = fromIntegral b
  approxInterv IntervInfiniteArbitrary = Exact IntervInfinite
  
  approxVec :: forall n .
               V.Vec n (Interval 'Arbitrary)
            -> Approximation (V.Vec n (Interval 'Standard))
  approxVec v =
    case findApproxIntervs stdVec of
      ([],_) -> Exact . fmap fromExact $ stdVec
      _      -> Bound Nothing (Just $ upperBound <$> stdVec)
    where
      stdVec :: V.Vec n (Approximation (Interval 'Standard))
      stdVec = fmap approxInterv v
  
      findApproxIntervs :: forall n' . V.Vec n' (Approximation (Interval 'Standard))
                        -> ([ Int ], [ Int ])
      findApproxIntervs v' = findApproxIntervs' 0 v' ([],[])
  
      findApproxIntervs' :: forall n' . Int
                         -> V.Vec n' (Approximation (Interval 'Standard))
                         -> ([ Int ], [ Int ])
                         -> ([ Int ], [ Int ])
      findApproxIntervs' _ V.Nil acc = acc
      findApproxIntervs' i (V.Cons x xs) (bixs, eixs) =
        findApproxIntervs' (i+1) xs $
          case x of
            Bound{} -> (i:bixs, eixs)
            Exact{} -> (bixs, i:eixs)
  
  instance Container (Interval 'Standard) where
    type MemberTyp (Interval 'Standard) = Int64
    type CompTyp (Interval 'Standard) = SInt64
  
    member 0 (IntervHoled _ _ b) = b
    member i (IntervHoled a b _) = i >= a && i <= b
    member _ _ = True
  
    compile (IntervHoled i1 i2 b) i
      | b = inRange i range
      | otherwise = inRange i range .&& i ./= 0
      where
        range = (fromIntegral i1, fromIntegral i2)
    compile IntervInfinite _ = sTrue
  
  instance Lattice (Interval 'Standard) where
    (IntervHoled lb ub noHole) \/ (IntervHoled lb' ub' noHole') =
      IntervHoled (min lb lb') (max ub ub') (noHole || noHole')
    _ \/ _ = top
    (IntervHoled lb ub noHole) /\ (IntervHoled lb' ub' noHole') =
      IntervHoled (max lb lb') (min ub ub') (noHole && noHole')
    int@IntervHoled{} /\ _ = int
    _ /\ int = int
  
  instance BoundedJoinSemiLattice (Interval 'Standard) where
    bottom = IntervHoled 0 0 False
  
  instance BoundedMeetSemiLattice (Interval 'Standard) where
    top = IntervInfinite
  
  -- instance BoundedLattice (Interval 'Standard)
  
  --------------------------------------------------------------------------------
  -- Union of cartesian products normal form
  --------------------------------------------------------------------------------
  
  type UnionNF n a = NE.NonEmpty (V.Vec n a)
  
  vecLength :: UnionNF n a -> V.Natural n
  vecLength = V.lengthN . NE.head
  
  instance Container a => Container (UnionNF n a) where
    type MemberTyp (UnionNF n a) = V.Vec n (MemberTyp a)
    type CompTyp (UnionNF n a) = V.Vec n (CompTyp a)
    member is = any (member' is)
      where
        member' is' space = and $ V.zipWith member is' space
  
    compile spaces is = foldr1 (.||) $ NE.map (`compile'` is) spaces
      where
        compile' space is' =
          foldr' (\(set, i) -> (.&&) $ compile set i) sTrue $ V.zip space is'
  
  instance BoundedMeetSemiLattice a => Lattice (UnionNF n a) where
    oi \/ oi' = oi <> oi'
    (/\) = CM.liftM2 (V.zipWith (/\))
  
  unfCompare :: forall a b n . ( Container a,          Container b
                               , MemberTyp a ~ Int64,  MemberTyp b ~ Int64
                               , CompTyp a ~ SInt64,   CompTyp b ~ SInt64
                               )
             => UnionNF n a -> UnionNF n b -> Ordering
  unfCompare oi oi' = unsafePerformIO $ do
      thmRes <- prove pred
      case thmRes of
        -- Tell the user if there was a hard proof error (e.g., if
        -- z3 is not installed/accessible).
        -- TODO: give more information
        ThmResult (ProofError _ msgs _) -> fail $ unlines msgs
        _ ->
          if modelExists thmRes
          then do
            ce <- counterExample thmRes
            case V.fromList ce of
               V.VecBox cev ->
                 case V.proveEqSize (NE.head oi) cev of
                   Just V.ReflEq ->
                     -- TODO: The second branch is defensive programming the
                     -- member check is not necessary unless the counter example
                     -- is bogus (it shouldn't be). Delete if it adversely
                     -- effects the performance.
                     if | cev `member` oi  -> return GT
                        | cev `member` oi' -> return LT
                        | otherwise -> fail
                           "Impossible: counter example is in \
                            \neither of the operands"
                   Nothing -> fail
                      "Impossible: Counter example size doesn't \
                      \match the original vector size."
          else return EQ
    where
      counterExample :: ThmResult -> IO [ Int64 ]
      counterExample thmRes =
        case getModelAssignment thmRes of
          Right (False, ce) -> return ce
          Right (True, _) -> fail "Returned probable model."
          Left str -> fail str
  
      pred :: Predicate
      pred = do
        freeVars <- (mkFreeVars . dimensionality) oi :: Symbolic [ SInt64 ]
        case V.fromList freeVars of
          V.VecBox freeVarVec ->
            case V.proveEqSize (NE.head oi) freeVarVec of
              Just V.ReflEq -> return $
                compile oi freeVarVec .== compile oi' freeVarVec
              Nothing -> fail $
                "Impossible: Free variables size doesn't match that of the " ++
                "union parameter."
      dimensionality = V.length . NE.head
  
  --------------------------------------------------------------------------------
  -- Optimise unions
  --------------------------------------------------------------------------------
  
  instance PO.PartialOrd Offsets where
    (Offsets s) <= (Offsets s') = s <= s'
    SetOfIntegers <= Offsets{} = False
    _ <= SetOfIntegers = True
  
  instance PO.PartialOrd (Interval 'Standard) where
    (IntervHoled lb ub p) <= (IntervHoled lb' ub' p') =
      (p' || not p) && lb >= lb' && ub <= ub'
    IntervInfinite <= IntervHoled{} = False
    _ <= IntervInfinite = True
  
  instance PO.PartialOrd a => PO.PartialOrd (V.Vec n a) where
    v <= v' = and $ V.zipWith (PO.<=) v v'
  
  optimise :: UnionNF n (Interval 'Standard) -> UnionNF n (Interval 'Standard)
  optimise = NE.fromList . maximas . fixedPointUnion . NE.toList
    where
      fixedPointUnion unf =
        let unf' = unionLemma . maximas $ unf
        in if unf' == unf then unf' else fixedPointUnion unf'
  
  sensibleGroupBy :: Eq a =>
                     (a -> a -> Ordering)
                  -> (a -> a -> Bool)
                  -> [ a ]
                  -> [ [ a ] ]
  sensibleGroupBy ord p l = nub . map (\el -> sortBy ord . filter (p el) $ l) $ l
  
  maximas :: [ V.Vec n (Interval 'Standard) ] -> [ V.Vec n (Interval 'Standard) ]
  maximas = nub
          . fmap (head . PO.maxima)
          . sensibleGroupBy ord (PO.<=)
    where
      ord a b = fromJust $ a `PO.compare` b
  
  -- | Union lemma says that if we have a product of intervals (as defined in
  -- the paper) and we union two that agrees in each dimension except one.
  -- The union is again a product of intervals that agrees with the original
  -- dimensions in all dimensions except the original differing one. At that
  -- point it is the union of intervals, which is itself still an interval.
  unionLemma :: [ V.Vec n (Interval 'Standard) ] -> [ V.Vec n (Interval 'Standard) ]
  unionLemma = map (foldr1 (V.zipWith (\/)))
             . sensibleGroupBy (\a b -> if a == b then EQ else LT) agreeButOne
    where
      -- This function returns true if two vectors agree at all points but one.
      -- It also holds if two vectors are identical.
      agreeButOne :: Eq a => V.Vec n a -> V.Vec n a -> Bool
      agreeButOne = go False
        where
          go :: Eq a => Bool -> V.Vec n a -> V.Vec n a -> Bool
          go _ V.Nil V.Nil = True
          go False (V.Cons x xs) (V.Cons y ys)
            | x == y = go False xs ys
            | otherwise = go True xs ys
          go True (V.Cons x xs) (V.Cons y ys)
            | x == y = go True xs ys
            | otherwise = False
  
  --------------------------------------------------------------------------------
  -- Injections for multiplicity and exactness
  --------------------------------------------------------------------------------
  
  data Approximation a = Exact a | Bound (Maybe a) (Maybe a)
    deriving (Eq, Show, Functor, Foldable, Traversable, Data, Typeable)
  
  fromExact :: Approximation a -> a
  fromExact (Exact a) = a
  fromExact _ = error "Can't retrieve from bounded as if it was exact."
  
  lowerBound :: Approximation a -> a
  lowerBound (Bound (Just a) _) = a
  lowerBound (Bound Nothing _) = error "Approximation doesn't have a lower bound."
  lowerBound (Exact a) = a
  
  upperBound :: Approximation a -> a
  upperBound (Bound _ (Just a)) = a
  upperBound (Bound _ Nothing) = error "Approximation doesn't have a upper bound."
  upperBound (Exact a) = a
  
  class Peelable a where
    type CoreTyp a
    peel :: a -> CoreTyp a
  
  data Multiplicity a = Mult a | Once a
    deriving (Eq, Show, Functor, Foldable, Traversable, Data, Typeable)
  
  instance Peelable (Multiplicity a) where
    type CoreTyp (Multiplicity a) = a
  
    peel (Mult a) = a
    peel (Once a) = a
  
  {-
  data Approximation a = Exact a | Lower a | Upper a
    deriving (Eq, Show, Functor, Data, Typeable)
  
  instance Peelable Approximation where
    peel (Exact a) = a
    peel (Lower a) = a
    peel (Upper a) = a
  -}
  