content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  {-# LANGUAGE QuasiQuotes, TypeApplications #-}
  
  module CompileBench (main) where
  
  import Control.Monad
  import Criterion.Measurement (secs)
  import Data.Text (Text)
  import qualified Data.Vector.Algorithms.Intro as V
  import qualified Data.Vector.Unboxed as V
  import Statistics.Sample (mean)
  import System.Mem
  import System.IO.Temp
  import System.FilePath ((</>))
  
  import Util.String.Quasi
  
  import Glean
  import Glean.Angle.Types (latestAngleVersion, AngleVersion(..))
  import Glean.Database.Open
  import Glean.Database.Test
  import Glean.Database.Types
  import Glean.Query.Thrift.Internal
  import qualified Glean.Schema.CodeCxx.Types as Code.Cxx
  import qualified Glean.Schema.Codemarkup.Types as Codemarkup
  import Glean.Types
  import Glean.Write.JSON (syncWriteJsonBatch)
  
  data Cfg = Cfg
    { cfgIters :: Int
    , cfgStop :: Double
    }
  
  type Description = String
  
  main :: IO ()
  main = do
    benchLargeDerivations
    benchTransformations
  
  benchTransformations :: IO ()
  benchTransformations =
    withSchemaFile latestAngleVersion schema $ \root file -> do
    let settings = [setRoot root, setSchemaPath file]
    repo <- withEmptyTestDB settings $ \env repo -> do
      void $ syncWriteJsonBatch env repo facts Nothing
      completeTestDB env repo
      return repo
    -- Open db for querying. We need to open the db again because schema
    -- evolutions and pruning are only triggered when the db is read-only.
    withTestEnv settings $ \env -> do
    bench env repo queries
    where
      queries =
        [ ("base predicate - no transformation"
          , getq $ angleData @Text $
            [s|
              S where
                # bind lhs                    # bind rhs
                X1 = x.P.2 {_ , _, _, _, _ }; X1 = x.P.2 {S1, _, _, _, _ };
                X2 = x.P.2 {S1, _, _, _, _ }; X2 = x.P.2 {S2, _, _, _, _ };
                X3 = x.P.2 {S2, _, _, _, _ }; X3 = x.P.2 {S3, _, _, _, _ };
                X4 = x.P.2 {S3, _, _, _, _ }; X4 = x.P.2 {S4, _, _, _, _ };
                X5 = x.P.2 {S4, _, _, _, _ }; X5 = x.P.2 {S5, _, _, _, _ };
                X6 = x.P.2 {S5, _, _, _, _ }; X6 = x.P.2 {S , _, _, _, _ };
            |]
          )
        , ("base predicate - with transformation"
          , getq $ angleData @Text $
            [s|
              S where
                # bind lhs          # bind rhs
                X1 = x.P.1 {_, _ }; X1 = x.P.1 {_, S1};
                X2 = x.P.1 {_, S1}; X2 = x.P.1 {_, S2};
                X3 = x.P.1 {_, S2}; X3 = x.P.1 {_, S3};
                X4 = x.P.1 {_, S3}; X4 = x.P.1 {_, S4};
                X5 = x.P.1 {_, S4}; X5 = x.P.1 {_, S5};
                X6 = x.P.1 {_, S5}; X6 = x.P.1 {_, S };
            |]
          )
        , ("derived predicate - no transformation"
          , getq $ angleData @Text $
            [s|
              S where x.Q.2 S =
                [ x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo"
                , x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo"
                , x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo"
                , x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo", x.Q.2 "foo"
                ][..];
            |]
          )
        , ("derived predicate - with transformation"
          , getq $ angleData @Text $
            [s|
              S where x.Q.1 S =
                [ x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo"
                , x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo"
                , x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo"
                , x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo", x.Q.1 "foo"
                ][..];
            |]
          )
        ]
      schema =
        [s|
          schema x.1 {
            predicate P : { x : nat, y : string }
            predicate Q : string
              S where
                # bind lhs      # bind rhs
                X1 = P {_, _ }; X1 = P {_, S1};
                X2 = P {_, S1}; X2 = P {_, S2};
                X3 = P {_, S2}; X3 = P {_, S3};
                X4 = P {_, S3}; X4 = P {_, S4};
                X5 = P {_, S4}; X5 = P {_, S5};
                X6 = P {_, S5}; X6 = P {_, S };
          }
          schema x.2 : x.1 {
            predicate P :
              { y : string
              , x : nat
              , a: maybe bool
              , b: maybe string
              , c: maybe string
              }
  
            predicate Q : string
              S where
                # bind lhs                # bind rhs
                X1 = P {_ , _, _, _, _ }; X1 = P {S1, _, _, _, _ };
                X2 = P {S1, _, _, _, _ }; X2 = P {S2, _, _, _, _ };
                X3 = P {S2, _, _, _, _ }; X3 = P {S3, _, _, _, _ };
                X4 = P {S3, _, _, _, _ }; X4 = P {S4, _, _, _, _ };
                X5 = P {S4, _, _, _, _ }; X5 = P {S5, _, _, _, _ };
                X6 = P {S5, _, _, _, _ }; X6 = P {S , _, _, _, _ };
          }
          schema x.2 evolves x.1
          schema all.1 : x.1, x.2 {}
        |]
  
      facts =
        [ mkBatch (PredicateRef "x.P" 2)
          [ [s|{ "key": { "x" : 1, "y": "A", "a": true, "b": "", "c": "" } }|] ]
        ]
  
      mkBatch ref facts = JsonFactBatch
        { jsonFactBatch_predicate = ref
        , jsonFactBatch_facts = facts
        , jsonFactBatch_unit = Nothing
        }
  
      withSchemaFile :: AngleVersion -> String -> (FilePath -> FilePath -> IO a) -> IO a
      withSchemaFile (AngleVersion version) str action = do
        withSystemTempDirectory "glean-dbtest" $ \root -> do
          let newSchemaFile = root </> "schema"
          appendFile newSchemaFile $ "version: " <> show version
          appendFile newSchemaFile str
          action root newSchemaFile
  
  benchLargeDerivations :: IO ()
  benchLargeDerivations =
    withEmptyTestDB [] $ \env repo -> do
    -- NB.Pruning and evolutions do not take place because the
    -- tests are run on a writable db.
    withOpenDatabase env repo $ void . return
    bench env repo queries
    where
      queries =
        [ ("codemarkup.search.earchByScope"
          , getq $ angleData @Code.Cxx.Entity $
          "codemarkup.search.SearchByScope { name = \"malloc\", scope = [], searchcase = Sensitive, kind = { just = Function }, entity = E }"
          )
        , ("codemarkup.EntityUses"
          , getq $ angleData @Codemarkup.EntityUses $
            [s|
                codemarkup.ResolveLocation {
                entity = E,
                location = {
                    location = { span = { length = 6, start = 1308 } },
                    file = "foo"
                }
                };
                codemarkup.EntityUses { target = E }
            |]
          )
        , ("FileEntityXRefLocations"
          , getq $ angleData @Codemarkup.FileEntityXRefLocations $
            [s|
              codemarkup.FileEntityXRefLocations { file = "foo" }
            |]
          )
        ]
  
  getq :: Query a -> UserQuery
  getq (Query q) = q
  
  bench :: Env -> Repo -> [(Description, UserQuery)] -> IO ()
  bench env repo queries =
    forM_ queries $ \(description,q) -> do
      xs <- measure (query env repo q) Cfg{cfgIters = 10, cfgStop = 1}
      let (compile, codegen) = V.unzip xs
      putStrLn $ " + " <> description
      putStrLn $ line "compile" compile
      putStrLn $ line "codegen" codegen
    where
      measure :: IO (Double,Double) -> Cfg -> IO (V.Vector (Double,Double))
      measure action = V.unfoldrM $ \Cfg{..} ->
        if cfgStop <= 0 && cfgIters <= 0
          then return Nothing
          else do
            performMinorGC
            (compile, codegen) <- action
            return $ Just ((compile, codegen), Cfg
              { cfgIters = max 0 $ cfgIters - 1
              , cfgStop = if compile >= cfgStop then 0 else cfgStop - compile
              })
  
      query env repo q = do
        UserQueryResults
          { userQueryResults_stats = Just UserQueryStats
              { userQueryStats_compile_time_ns = Just compile
              , userQueryStats_codegen_time_ns = Just codegen }
          } <- userQuery env repo q
        return (fromNS compile, fromNS codegen)
        where
          fromNS n = fromIntegral n / 1000000000
  
      line :: String -> V.Vector Double -> String
      line h xs = replicate 6 ' ' ++ h ++ " " ++ display xs
        where
  
      display :: V.Vector Double -> String
      display xs = unwords
        [field "min" $ secs $ V.head ts
        ,field "p50" $ secs $ ts V.! (V.length ts `div` 2 + 1)
        ,field "avg" $ secs $ mean ts
        ,field "max" $ secs $ V.last ts
        ,field "its" $ show $ V.length xs
        ]
        where
          ts = V.modify V.sort xs
          field s t = s ++ ": " ++ t
  