-- Copyright 2022 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

module OccurrenceSpec (spec) where

import Prelude hiding ((.), seq)

import Control.Monad.Reader.Class
import Data.Map.Strict qualified as M
import Data.String
import Test.Hspec

import MTL1
import IRVariants
import Name hiding (NameMap)
import Occurrence

-- === Embedded DSL for spelling example Access objects ===

data UAccessOpen =
    ULoc [UIxExpr] Count
  | UAny [UAccessOpen]
  | UAll [UAccessOpen]
  | UForEach String UAccessOpen

data UAccess = Free [String] UAccessOpen

closed :: UAccessOpen -> UAccess
closed = Free []

data UIxExpr =
    UVar String
  | UProduct [UIxExpr]
  | UInject Int UIxExpr
  | UDeterministic [String]

instance IsString UIxExpr where
  fromString = UVar

resolveNames :: UAccess -> Abs (Nest (NameBinder ('AtomNameC SimpIR))) Access 'VoidS
resolveNames (Free names uaccess) = runResolverM
  $ withFreshNamesM (map getNameHint names) \bs ->
    local ((NameMap $ M.fromList
            $ zip names (nestToList (\b -> sink $ binderName b) bs)) <>) do
      access <- resolveNames' uaccess
      return $ Abs bs access

withFreshNamesM
  :: (Color c, ScopeExtender m, Distinct n)
  => [NameHint]
  -> (forall l. DExt n l => (Nest (NameBinder c)) n l -> m l a)
  -> m n a
withFreshNamesM [] cont = cont Empty
withFreshNamesM (hint:hints) cont = withFreshM hint \b ->
  withFreshNamesM hints \bs -> cont $ Nest b bs

newtype NameMap (n::S) = NameMap (M.Map String (Name ('AtomNameC SimpIR) n))
  deriving (Semigroup, Monoid)

instance SinkableE NameMap where
  sinkingProofE c (NameMap m) = NameMap $ fmap (sinkingProofE c) m

newtype ResolverM (n::S) a = ResolverM {
  runResolverM' :: ReaderT1 NameMap ScopeReaderM n a }
  deriving (Functor, Applicative, Monad,
            ScopeReader, ScopeExtender, MonadReader (NameMap n))

runResolverM :: ResolverM 'VoidS a -> a
runResolverM action = runScopeReaderM emptyOutMap
  $ runReaderT1 (NameMap M.empty)
  $ runResolverM' action

resolveNames' :: UAccessOpen -> ResolverM n (Access n)
resolveNames' (ULoc ixs ct) = flip Location ct <$> mapM resolveNames'' ixs
resolveNames' (UAny uaccesses) = Any <$> mapM resolveNames' uaccesses
resolveNames' (UAll uaccesses) = All <$> mapM resolveNames' uaccesses
resolveNames' (UForEach str uaccess) =
  withFreshM (getNameHint str) \b ->
    ForEach b <$> (local ((NameMap $ M.singleton str (binderName b)) <>)
      $ resolveNames' uaccess)

resolveNames'' :: UIxExpr -> ResolverM n (IxExpr n)
resolveNames'' (UVar str) = Var <$> lookupName str
resolveNames'' (UProduct elts) = Product <$> mapM resolveNames'' elts
resolveNames'' (UInject i elt) = Inject i <$> resolveNames'' elt
resolveNames'' (UDeterministic names) = Deterministic <$> mapM lookupName names

lookupName :: String -> ResolverM n (Name ('AtomNameC SimpIR) n)
lookupName str = do
  (NameMap names) <- ask
  case M.lookup str names of
    Just n -> return n
    Nothing -> error $ "Unknown name " ++ str

for :: String -> UAccessOpen -> UAccessOpen
for = UForEach

xs :: UAccessOpen
xs = ULoc [] One

(.) :: UAccessOpen -> UIxExpr -> UAccessOpen
(ULoc items ct) . item = ULoc (items ++ [item]) ct
_ . _ = error "Invalid indexing expression"

seq :: [UAccessOpen] -> UAccessOpen
seq = UAll

tup :: [UIxExpr] -> UIxExpr
tup = UProduct

alt :: [UAccessOpen] -> UAccessOpen
alt = UAny

inj :: Int -> UIxExpr -> UIxExpr
inj = UInject

left :: UIxExpr -> UIxExpr
left = inj 0

right :: UIxExpr -> UIxExpr
right = inj 1

det :: [String] -> UIxExpr
det = UDeterministic

free :: [String] -> UAccessOpen -> UAccess
free = Free

-- === The actual spec ===

answerC :: UAccessOpen -> DynUseInfo
answerC open = answer $ closed open

answer :: UAccess -> DynUseInfo
answer uaccess = case resolveNames uaccess of
  (Abs _ open) -> approxConst $ collapse $ interp $ open

spec :: Spec
spec = do
  describe "Occurrence computes the right answer on" do
    it "just an access" do
      answerC xs `shouldBe` (0, One)
    it "1-D identity array" do
      answerC (for "i" $ xs . "i") `shouldBe` (1, One)
    it "two sequential accesses" do
      answerC (seq [for "i" $ xs . "i", for "i" $ xs . "i"])
        `shouldBe` (1, Bounded 2)
    it "in-body sequential access" do
      answerC (for "i" $ seq [xs . "i", xs . "i"]) `shouldBe` (1, Bounded 2)
    it "two alternate accesses" do
      answerC (alt [for "i" $ xs . "i", for "i" $ xs . "i"])
        `shouldBe` (1, One)
    it "in-body access alternation" do
      answerC (for "i" $ alt [xs . "i", xs . "i"]) `shouldBe` (1, One)
    it "2-D identity array" do
      answerC (for "i" $ for "j" $ xs . "i" . "j") `shouldBe` (2, One)
    it "2-D transposition" do
      answerC (for "i" $ for "j" $ xs . "j" . "i") `shouldBe` (2, One)
    it "a known function of the index" do
      answerC (for "i" $ xs . left "i") `shouldBe` (1, One)
    it "separately accessing two variants of the index type" do
      answerC (seq [ for "i" $ xs . left  "i"
                   , for "i" $ xs . right "i"]) `shouldBe` (1, One)
    it "unknown indexing in a loop" do
      answerC (for "i" $ xs . det ["i"]) `shouldBe` (1, Unbounded)
    it "loop-invariant access" do
      answerC (for "i" $ xs) `shouldBe` (0, Unbounded)
    it "2-D loop-invariant access" do
      answerC (for "i" $ for "j" $ xs . "j") `shouldBe` (1, Unbounded)
    it "a simple access with a ternary sum index type" do
      answerC (for "i" $ seq [ xs . inj 1 "i"
                             , xs . inj 2 "i"
                             , xs . inj 3 "i"
                             ]) `shouldBe` (1, One)
    it "taking the trace" do
      -- TODO Actually, xs is safe to inline even if we can only see
      -- one dimension, because the first index proves unique access
      -- without requiring the second one.  So this answer should
      -- really be (1, One); but the analysis doesn't do that yet.
      answerC (for "i" $ xs . "i" . "i") `shouldBe` (2, One)
    it "sequencing distinct alternate accesses" do
      let access = for "i" $
            seq [ alt [xs . inj 1 "i", xs . inj 2 "i"]
                , alt [xs . inj 3 "i", xs . inj 4 "i"]
                ]
      answerC access `shouldBe` (1, One)
    it "sequencing colliding alternate accesses" do
      let access = for "i" $
            seq [ alt [xs . inj 1 "i", xs . inj 2 "i"]
                , alt [xs . inj 1 "i", xs . inj 4 "i"]
                ]
      answerC access `shouldBe` (1, Bounded 2)
    it "sequencing non-colliding alternate accesses" do
      let access = for "i" $
            seq [ alt [xs . inj 1 "i", xs . inj 1 "i"]
                , alt [xs . inj 3 "i", xs . inj 4 "i"]
                ]
      answerC access `shouldBe` (1, Bounded 1)
    it "taking a 'trace' with one unknown index" do
      answerC (for "i" $ xs . det ["i"] . "i") `shouldBe` (2, One)
    it "a tuple-typed index" do
      answerC (for "i" $ for "j" $ xs . tup ["i", "j"]) `shouldBe` (1, One)
    it "a tuple-typed index with an Either inside" do
      answerC (seq [ for "i" $ for "j" $ xs . tup ["i", left "j"]
                   , for "i" $ for "j" $ xs . tup ["i", right "j"]])
        `shouldBe` (1, One)
  describe "Occurrence computes the right answer on open terms" do
    it "accessing a free index" do
      answer (free ["i"] $ xs . "i") `shouldBe` (1, One)
    it "accessing a free index in a loop" do
      -- TODO Same conservativeness here: one dimension is enough
      -- to prove safety
      answer (free ["j"] $ for "i" $ xs . "i" . "j") `shouldBe` (2, One)
    it "accessing a free index in a loop transposed" do
      -- Here, on the other hand, we really need two dimensions to
      -- inline without duplicating work.
      answer (free ["j"] $ for "i" $ xs . "j" . "i") `shouldBe` (2, One)
    it "accessing potentially colliding free indices" do
      -- Here we can't be completely safe inlining, because j and k
      -- could collide and force us to duplicate work on that one
      -- element in that one circumstance.  However, a more
      -- sophisticated analysis could discover that it may be worth
      -- either (1) just eating the duplication, if the j and k spaces
      -- are large enough and the body of xs is uniform enough, and/or
      -- (2) adding a guard condition for `j == k` and inlining into
      -- the branch where this is false.
      answer (free ["j", "k"] $ seq
               [ for "i" $ xs . "i" . "j"
               , for "i" $ xs . "i" . "k"
               ]) `shouldBe` (2, Bounded 2)
    it "accessing unknown free index" do
      -- Again, the real answer is (0, One), but the analysis doesn't
      -- know that yet.
      answer (free ["i"] $ xs . (det ["i"])) `shouldBe` (1, One)

-- === TODO Property-based testing ===
--
-- The Occurrence data structures should obey some invariants:
-- - Each MaxPlus instance should obey the Semigroup laws (except
--   annihilation)
-- - max(answer(a), answer(b)) == answer(max(a, b))
-- - sum(answer(a), answer(b)) >= answer(sum(a, b))
--
-- These would be good targets for property-based testing.  The open
-- issue is generating well-typed synthetic Access, Use, or
-- CollapsedUse terms.
--
-- One way to do that would be to
-- - Generate a type for a binding -- an array, indexed by a product
--   or sum-typed index, with given elements or alternatives, etc.
-- - Generate an Access for a binding of that type
--   - Generate some free index names
--   - When generating a `for`, generate a new index name
--   - When generating a Location, generate expressions according
--     to the binding's type
--
-- Issue: What type to make the generated iteration-space binders?
-- Make something up and hope expressions become possible?  Reuse
-- (parts of) the type of the binding?  Make a type variable and infer
-- it based on leaf Location expressions?