{-
  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 #-}
{-# LANGUAGE TypeApplications #-}
module QueryBench (main) where

import Control.Concurrent.Async
import Control.Monad (void)
import Criterion.Types
import Data.Default

import Util.String.Quasi

import Glean
import Glean.Angle as Angle
import qualified Glean.Schema.CodeCxx.Types as Code.Cxx
import qualified Glean.Schema.Codemarkup.Types as Codemarkup
import qualified Glean.Schema.Cxx1.Types as Cxx
import qualified Glean.Schema.GleanTest.Types as Glean.Test
import qualified Glean.Schema.Sys.Types as Sys
import Glean.Util.Benchmark

import BenchDB

main :: IO ()
main = benchmarkMain $ \run -> withBenchDB 10000 $ \env repo -> do
  let
    nestedAngle :: Query Cxx.FunctionName
    nestedAngle = angle "cxx1.FunctionName { name = \"x1\" }"

    pageAngle :: Query Cxx.Name
    pageAngle = limit 1 $ angle "cxx1.Name \"y\".."

    pageNestedAngle :: Query Cxx.FunctionName
    pageNestedAngle = limit 1 $ angle "cxx1.FunctionName { name = \"y\".. }"

    simpleAngle :: Query Sys.Blob
    simpleAngle = angle "sys.Blob \"x1\""

    complexQueryAngle :: Query Glean.Test.Predicate
    complexQueryAngle = Angle.query $ predicate @Glean.Test.Predicate $ rec $
      field @"string_" (string "x1") end

    -- test compiling a complex derived predicate, should let us
    -- spot regressions in the compiler pipeline.
    compile :: Query Code.Cxx.Entity
    compile = angleData @Code.Cxx.Entity $
      "E where codemarkup.search.SearchByScope { name = \"malloc\", scope = [], searchcase = Sensitive, kind = { just = Function }, entity = E }"

    compile2 :: Query Codemarkup.EntityUses
    compile2 = angleData @Codemarkup.EntityUses $
      [s|
        codemarkup.ResolveLocation {
          entity = E,
          location = {
            location = { span = { length = 6, start = 1308 } },
            file = "foo"
          }
        };
        codemarkup.EntityUses { target = E }
      |]

    compile3 :: Query Codemarkup.FileEntityXRefLocations
    compile3 = angleData @Codemarkup.FileEntityXRefLocations $
      [s|
        codemarkup.FileEntityXRefLocations { file = "foo" }
      |]

    arrayPrefix :: Query Nat
    arrayPrefix = angleData @Nat
      [s| N where
         glean.test.Predicate { array_of_nat = [N, ..]}
       |]

  let stacked = Repo "test" "stacked"
  let
    dependencies =
      Dependencies_stacked $ Stacked name hash Nothing
      where (Repo name hash) = repo
  void $ kickOffDatabase env def
    { kickOff_repo = stacked
    , kickOff_dependencies = Just dependencies
    }
  run
      [ bench "complex-parallel" $ whnfIO
        $ mapConcurrently (\io -> io >>= \x -> x `seq` return x)
        $ replicate 1000
        $ runQuery env repo complexQueryAngle
      , bench "simple" $ whnfIO $
          runQuery env repo simpleAngle
      , bench "nested" $ whnfIO $
          runQuery env repo nestedAngle
      , bench "page" $ whnfIO $
          runQuery_ env repo pageAngle
      , bench "pagenested" $ whnfIO $
          runQuery_ env repo pageNestedAngle
      , bench "compile" $ whnfIO $
          runQuery_ env repo compile
      , bench "compile2" $ whnfIO $
          runQuery_ env repo compile2
      , bench "compile3" $ whnfIO $
          runQuery_ env repo compile3
      , bench "array_prefix" $ whnfIO $
          runQuery_ env repo arrayPrefix
      , bgroup "stacked"
        [ bench "page" $ whnfIO $
            runQuery_ env repo pageAngle
        , bench "pagenested" $ whnfIO $
            runQuery_ env repo pageNestedAngle
        ]
      ]