content: {-# LANGUAGE DataKinds, TypeFamilies, DeriveGeneric, Arrows #-}
  
  module Test.Language.Souffle.AnalysisSpec
    ( module Test.Language.Souffle.AnalysisSpec
    ) where
  
  import Prelude hiding ((.), id)
  import Control.Arrow
  import Control.Category
  import Test.Hspec
  import Data.Profunctor
  import GHC.Generics
  import Control.Monad.IO.Class
  import Language.Souffle.Analysis
  import qualified Language.Souffle.Interpreted as Souffle
  
  data Path = Path
  
  data Edge = Edge String String
    deriving stock (Eq, Show, Generic)
  
  data Reachable = Reachable String String
    deriving stock (Eq, Show, Generic)
  
  instance Souffle.Program Path where
    type ProgramFacts Path = '[Edge, Reachable]
    programName = const "path"
  
  instance Souffle.Fact Edge where
    type FactDirection Edge = 'Souffle.InputOutput
    factName = const "edge"
  
  instance Souffle.Fact Reachable where
    type FactDirection Reachable = 'Souffle.Output
    factName = const "reachable"
  
  instance Souffle.Marshal Edge
  instance Souffle.Marshal Reachable
  
  data Results = Results [Reachable] [Edge]
    deriving stock (Eq, Show)
  
  pathAnalysis :: Souffle.Handle Path
               -> Analysis Souffle.SouffleM [Edge] [Reachable]
  pathAnalysis h =
    mkAnalysis (Souffle.addFacts h) (Souffle.run h) (Souffle.getFacts h)
  
  -- A little bit silly, but good enough to test different forms of application with
  pathAnalysis' :: Souffle.Handle Path
               -> Analysis Souffle.SouffleM [Edge] [Edge]
  pathAnalysis' h =
    mkAnalysis (Souffle.addFacts h) (Souffle.run h) (Souffle.getFacts h)
  
  data RoundTrip = RoundTrip
  
  newtype StringFact = StringFact String
    deriving stock (Eq, Show, Generic)
  
  instance Souffle.Program RoundTrip where
    type ProgramFacts RoundTrip = '[StringFact]
  
    programName = const "round_trip"
  
  instance Souffle.Fact StringFact where
    type FactDirection StringFact = 'Souffle.InputOutput
  
    factName = const "string_fact"
  
  instance Souffle.Marshal StringFact
  
  
  roundTripAnalysis :: Souffle.Handle RoundTrip
                    -> Analysis Souffle.SouffleM [Reachable] [StringFact]
  roundTripAnalysis h =
    mkAnalysis addFacts (Souffle.run h) (Souffle.getFacts h)
    where
      addFacts rs = do
        Souffle.addFacts h $ map (\(Reachable a _) -> StringFact a) rs
  
  withSouffle :: Souffle.Program a => a -> (Souffle.Handle a -> Souffle.SouffleM ()) -> IO ()
  withSouffle prog f = Souffle.runSouffle prog $ \case
    Nothing -> error "Failed to load program"
    Just h -> f h
  
  edges :: [Edge]
  edges = [Edge "a" "b", Edge "b" "c", Edge "b" "d", Edge "d" "e"]
  
  spec :: Spec
  spec = describe "composing analyses" $ parallel $ do
    it "supports fmap" $ do
      withSouffle Path $ \h -> do
        let analysis = pathAnalysis h
            analysis' = fmap length analysis
        count <- execAnalysis analysis' edges
        liftIO $ count `shouldBe` 8
  
    describe "analysis used as a profunctor" $ parallel $ do
      it "supports lmap" $ do
        withSouffle Path $ \h -> do
          let inputs = [("a", "b"), ("b", "c")]
              analysis = pathAnalysis h
              analysis' = lmap (map (uncurry Edge)) analysis
          rs <- execAnalysis analysis' inputs
          liftIO $ rs `shouldBe` [ Reachable "a" "b"
                                 , Reachable "a" "c"
                                 , Reachable "b" "c"
                                 ]
  
      it "supports rmap" $ do
        withSouffle Path $ \h -> do
          let analysis = pathAnalysis h
              analysis' = rmap length analysis
          count <- execAnalysis analysis' edges
          liftIO $ count `shouldBe` 8
  
    it "supports applicative composition" $
      withSouffle Path $ \hPath -> do
        let analysis1 = pathAnalysis hPath
            analysis2 = pathAnalysis' hPath
            analysis = Results <$> analysis1 <*> analysis2
            inputs = [Edge "a" "b", Edge "b" "c"]
            reachables = [ Reachable "a" "b"
                         , Reachable "a" "c"
                         , Reachable "b" "c"
                         ]
        results <- execAnalysis analysis inputs
        liftIO $ results `shouldBe` Results reachables inputs
  
    it "supports semigroupal composition" $ do
      withSouffle Path $ \h -> do
        let analysis = pathAnalysis h
            analysis' = analysis <> analysis
        rs <- execAnalysis analysis' [Edge "a" "b", Edge "b" "c"]
        let results = [ Reachable "a" "b"
                      , Reachable "a" "c"
                      , Reachable "b" "c"
                      ]
            results' = mconcat $ replicate 2 results
        liftIO $ rs `shouldBe` results'
  
    it "supports mempty" $ do
      withSouffle Path $ \_ -> do
        let analysis :: Analysis Souffle.SouffleM [Edge] [Reachable]
            analysis = mempty
        rs <- execAnalysis analysis [Edge "a" "b", Edge "b" "c"]
        liftIO $ rs `shouldBe` []
  
    it "supports converting an analysis to a monadic function" $ do
      withSouffle Path $ \h -> do
        let analysis = pathAnalysis h
        rs <- execAnalysis analysis [Edge "a" "b", Edge "b" "c"]
        let results = [ Reachable "a" "b"
                      , Reachable "a" "c"
                      , Reachable "b" "c"
                      ]
        liftIO $ rs `shouldBe` results
  
    describe "analysis used as a category" $ parallel $ do
      it "supports 'id'" $ do
        withSouffle Path $ \_ -> do
          let analysis :: Analysis Souffle.SouffleM [Edge] [Edge]
              analysis = id
          edges' <- execAnalysis analysis edges
          liftIO $ edges' `shouldBe` edges
  
      it "supports sequential composition using (.)" $ do
        withSouffle Path $ \h -> do
          let reachableToFlippedEdge (Reachable a b) = Edge b a
              analysis1 = pathAnalysis h
              analysis2 = lmap (map reachableToFlippedEdge) $ pathAnalysis h
          rs <- execAnalysis (analysis2 . analysis1) [Edge "a" "b", Edge "b" "c"]
          let results = [ Reachable "a" "a"
                        , Reachable "a" "b"
                        , Reachable "a" "c"
                        , Reachable "b" "a"
                        , Reachable "b" "b"
                        , Reachable "b" "c"
                        , Reachable "c" "a"
                        , Reachable "c" "b"
                        , Reachable "c" "c"
                        ]
          liftIO $ rs `shouldBe` results
  
    describe "analysis used as an arrow" $ parallel $ do
      it "supports 'arr'" $ do
        withSouffle Path $ \_ -> do
          let analysis :: Analysis Souffle.SouffleM Int Int
              analysis = arr (+1)
          result1 <- execAnalysis analysis 41
          result2 <- execAnalysis (arr id) 41
          liftIO $ result1 `shouldBe` 42
          liftIO $ result2 `shouldBe` 41
  
      it "supports 'first'" $ do
        withSouffle Path $ \_ -> do
          let analysis :: Analysis Souffle.SouffleM (Int, Bool) (Int, Bool)
              analysis = first (arr (+1))
              input = (41, True)
          result <- execAnalysis analysis input
          liftIO $ result `shouldBe` (42, True)
  
      it "supports 'second'" $ do
        withSouffle Path $ \_ -> do
          let analysis :: Analysis Souffle.SouffleM (Bool, Int) (Bool, Int)
              analysis = second (arr (+1))
              input = (True, 41)
          result <- execAnalysis analysis input
          liftIO $ result `shouldBe` (True, 42)
  
      it "supports (***)" $ do
        withSouffle Path $ \_ -> do
          let analysis :: Analysis Souffle.SouffleM (Bool, Int) (Bool, Int)
              analysis = arr not *** arr (+1)
              input = (True, 41)
          result <- execAnalysis analysis input
          liftIO $ result `shouldBe` (False, 42)
  
      it "supports (&&&)" $ do
        withSouffle Path $ \_ -> do
          let analysis :: Analysis Souffle.SouffleM Int (Bool, Int)
              analysis = arr (== 1000) &&& arr (+1)
              input = 41
          result <- execAnalysis analysis input
          liftIO $ result `shouldBe` (False, 42)
  
      it "supports arrow notation" $ do
        withSouffle Path $ \h -> do
          liftIO $ withSouffle RoundTrip $ \h' -> do
            let arrowAnalysis = proc es -> do
                  rs <- pathAnalysis h -< es
                  strs <- roundTripAnalysis h' -< rs
                  returnA -< strs
            result <- execAnalysis arrowAnalysis edges
            liftIO $ result `shouldBe` [ StringFact "a"
                                       , StringFact "b"
                                       , StringFact "d"
                                       ]
  
      it "supports case expressions in arrow notation" $ do
        withSouffle Path $ \h -> do
          let analysis =  proc es -> do
                rs <- pathAnalysis h -< es
                case rs of
                  [] -> returnA -< []
                  rs' -> returnA -< take 2 rs'
          result <- execAnalysis analysis edges
          let expected = [ Reachable "a" "b", Reachable "a" "c" ]
          liftIO $ result `shouldBe` expected
  
  