Count.hs

{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}

module App.Commands.Count
  ( cmdCount
  ) where

import Control.Lens
import Control.Monad
import Data.Generics.Product.Any
import Data.Semigroup                                 ((<>))
import Data.Word
import HaskellWorks.Data.FromForeignRegion
import HaskellWorks.Data.Json.LightJson
import HaskellWorks.Data.Json.Query
import HaskellWorks.Data.Json.Standard.Cursor.Generic
import HaskellWorks.Data.MQuery
import HaskellWorks.Data.MQuery.Micro
import HaskellWorks.Data.RankSelect.CsPoppy1
import HaskellWorks.Data.TreeCursor
import Options.Applicative                            hiding (columns)

import qualified App.Commands.Types                                   as Z
import qualified Data.ByteString                                      as BS
import qualified Data.ByteString.Internal                             as BSI
import qualified Data.DList                                           as DL
import qualified Data.Vector.Storable                                 as DVS
import qualified HaskellWorks.Data.BalancedParens.RangeMin            as RM
import qualified HaskellWorks.Data.Json.Standard.Cursor.Fast          as JCF
import qualified HaskellWorks.Data.Json.Standard.Cursor.Internal.IbBp as IBBP
import qualified HaskellWorks.Data.RankSelect.CsPoppy.Internal.Alpha1 as A1
import qualified System.IO                                            as IO
import qualified System.IO.MMap                                       as IO

siblings :: GenericCursor BSI.ByteString CsPoppy1 (RM.RangeMin CsPoppy1) -> [GenericCursor BSI.ByteString CsPoppy1 (RM.RangeMin CsPoppy1)]
siblings c = c:cs
  where cs = case nextSibling c of
          Just d  -> siblings d
          Nothing -> []

runCount :: Z.CountOptions -> IO ()
runCount opts = do
  let inputFile   = opts ^. the @"inputFile"
  let expression  = opts ^. the @"expression"
  jsonFr    <- IO.mmapFileForeignPtr inputFile IO.ReadOnly Nothing
  let jsonBs  = fromForeignRegion jsonFr    :: BS.ByteString

  cursor <- case opts ^. the @"indexes" of
    Just indexes -> do
      let ibIndex     = indexes ^. the @"ibIndex"
      let bpIndex     = indexes ^. the @"bpIndex"
      jsonIbFr  <- IO.mmapFileForeignPtr ibIndex   IO.ReadOnly Nothing
      jsonBpFr  <- IO.mmapFileForeignPtr bpIndex   IO.ReadOnly Nothing
      let jsonIb  = fromForeignRegion jsonIbFr  :: DVS.Vector Word64
      let jsonBp  = fromForeignRegion jsonBpFr  :: DVS.Vector Word64

      return $ GenericCursor jsonBs (makeCsPoppy jsonIb) (RM.mkRangeMin (makeCsPoppy jsonBp)) 1
    Nothing -> do
      let !ibip = JCF.simdToIbBp jsonBs
      let !c    = JCF.fromBsIbBp jsonBs ibip
      return c

  let q = MQuery (DL.fromList $ fmap lightJsonAt (siblings cursor))

  putPretty $ q >>= (entry >=> named expression) & count

  return ()

optsFileIndex :: Parser Z.FileIndexes
optsFileIndex = Z.FileIndexes
  <$> strOption
    (   long "ib-index"
    <>  help "IB index"
    <>  metavar "FILE"
    )
  <*> strOption
    (   long "bp-index"
    <>  help "BP index"
    <>  metavar "FILE"
    )

optsCount :: Parser Z.CountOptions
optsCount = Z.CountOptions
  <$> strOption
        (   long "input"
        <>  short 'i'
        <>  help "Input JSON file"
        <>  metavar "FILE"
        )
  <*> optional optsFileIndex
  <*> option auto
        (   long "expression"
        <>  help "JSON expression"
        <>  metavar "EXPRESSION"
        )

cmdCount :: Mod CommandFields (IO ())
cmdCount = command "count"  $ flip info idm $ runCount <$> optsCount