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 TypeApplications #-}
  module BenchDB (withBenchDB) where
  
  import Control.Monad
  import qualified Data.ByteString.Char8 as BC
  import Data.Default
  import qualified Data.Text as Text
  
  import Glean
  import Glean.Backend.Types (loadPredicates)
  import Glean.Database.Test (withEmptyTestDB, completeTestDB)
  import Glean.Database.Open
  import Glean.Database.Write.Batch
  import qualified Glean.Schema.Sys.Types as Sys
  import qualified Glean.Schema.Cxx1.Types as Cxx
  import qualified Glean.Schema.GleanTest.Types as Glean.Test
  import qualified Glean.Schema.Sys as Sys
  import qualified Glean.Schema.Cxx1 as Cxx
  import qualified Glean.Schema.GleanTest as Glean.Test
  import Glean.Typed
  
  withBenchDB :: Int -> (forall b . Backend b => b -> Repo -> IO a) -> IO a
  withBenchDB num act = withEmptyTestDB [] $ \env repo -> do
    withOpenDatabase env repo $ \odb ->
      void $ return odb
  
    predicates <- loadPredicates env repo
      [ Sys.allPredicates
      , Cxx.allPredicates
      , Glean.Test.allPredicates
      ]
    batch <- buildBatch predicates Nothing $ do
      blobs <- mapM (makeFact @Sys.Blob)
        [ "x" <> BC.pack (show n) | n <- [1..num] ]
  
      names <- mapM (makeFact @Cxx.Name)
        [ "x" <> Text.pack (show n) | n <- [1..num] ]
      mapM_ (makeFact @Cxx.FunctionName . Cxx.FunctionName_key_name) names
  
      mapM_ (makeFact @Glean.Test.Predicate)
        [ def
          { Glean.Test.kitchenSink_string_ = "x" <> Text.pack (show n)
          , Glean.Test.kitchenSink_pred = blob
          , Glean.Test.kitchenSink_sum_ = Glean.Test.KitchenSink_sum__d blob
          , Glean.Test.kitchenSink_named_record_ =
              def { Glean.Test.rec_beta = Glean.Test.Sum_wed True }
          , Glean.Test.kitchenSink_named_sum_ =
              Glean.Test.Sum_wed True
          , Glean.Test.kitchenSink_maybe_ = Nothing
          , Glean.Test.kitchenSink_array_of_nat = map Nat [1 .. fromIntegral n]
          }
  
        | (n,blob) <- zip [(1::Int)..] blobs
        ]
  
      ynames <- mapM (makeFact @Cxx.Name)
        [ "y" <> Text.pack (show n) | n <- [1..1000::Int] ]
      mapM_ (makeFact @Cxx.FunctionName . Cxx.FunctionName_key_name) ynames
  
    void $ syncWriteDatabase env repo batch
    completeTestDB env repo
    act env repo
  