{-
  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 RecursiveDo #-}
module Glean.RTS.Traverse
  ( CompiledTraversal
  , genTraversal
  ) where

import Control.Monad

import Glean.Bytecode.Types
import Glean.RTS.Foreign.Bytecode
import Glean.RTS.Types
import Glean.RTS.Bytecode.Code
import Glean.RTS.Bytecode.Gen.Issue
import GHC.Stack (HasCallStack)

-- | Type tag for Subroutine
data CompiledTraversal

traversal
  :: HasCallStack => Register ('Fun '[ 'Word, 'Word ])
  -> Register 'DataPtr
  -> Register 'DataPtr
  -> Type
  -> Code ()
traversal callback input inputend ty = go False (repType ty)
  where
    -- if refs is True, we *must* leave the input pointer pointing
    -- after the value, because we're going to traverse more data.
    --
    -- If refs is False, we *must not* traverse the value unless
    -- (hasRefs ty) is True, because we might not have traversed the
    -- previous value.
    go refs ByteRep = when refs $ addConst 1 input
    go refs NatRep = when refs $ inputSkipNat input inputend
    go refs StringRep = when refs $ inputSkipTrustedString input inputend
    go refs (ArrayRep elty) = when (refs || hasRefs elty) $ local $ \size -> do
      inputNat input inputend size
      case elty of
        ByteRep -> add size input
        _ -> mdo
          jumpIf0 size end
          loop <- label
          go True elty
          decrAndJumpIfNot0 size loop
          end <- label
          return ()
    go refs (TupleRep tys) = mapM_ (uncurry go) $ zip
      (drop 1 $ scanr (\ty refs -> hasRefs ty || refs) refs tys)
      tys
    go refs (SumRep tys)
      | all isUnit tys = when refs $ if length tys <= 127
          then addConst 1 input
          else inputSkipNat input inputend
      | otherwise = when (refs || any hasRefs tys) $ mdo
          local $ \sel -> do
            -- TODO: read byte rather than Nat if arity <= 127
            inputNat input inputend sel
            let unknownAlt = [end]
            select sel (alts ++ unknownAlt)
          raise "selector out of range"
          alts <- forM tys $ \ty -> do
            alt <- label
            go refs ty
            jump end
            return alt
          end <- label
          return ()
    go refs (SetRep elty) = when (refs || hasRefs elty) $ local $ \size -> mdo
      inputNat input inputend size
      jumpIf0 size end
      loop <- label
      go True elty
      decrAndJumpIfNot0 size loop
      end <- label
      return ()
    go _ (PredicateRep (Pid pid)) = local $ \ide -> do
      inputNat input inputend ide
      pidr <- constant (fromIntegral pid)
      callFun_2_0 callback ide pidr

-- | Generate a subroutine which traverses a clause (fact key + value)
-- and invokes the supplied callback function for each fact ID
-- contained in it.
--
-- NOTE: The subroutine assumes that the clause is type correct.
genTraversal :: HasCallStack => Type -> Type -> IO (Subroutine CompiledTraversal)
genTraversal key_ty val_ty =
  fmap snd $ generate Optimised $ \handler clause_begin key_end clause_end -> do
    traversal handler clause_begin key_end key_ty
    traversal handler key_end clause_end val_ty
    ret

hasRefs :: Rep a -> Bool
hasRefs ByteRep = False
hasRefs NatRep = False
hasRefs StringRep = False
hasRefs (ArrayRep elty) = hasRefs elty
hasRefs (TupleRep tys) = any hasRefs tys
hasRefs (SumRep tys) = any hasRefs tys
hasRefs (SetRep elty) = hasRefs elty
hasRefs PredicateRep{} = True

isUnit :: Rep a -> Bool
isUnit ByteRep = False
isUnit NatRep = False
isUnit StringRep = False
isUnit ArrayRep{} = False
isUnit (TupleRep tys) = all isUnit tys
isUnit SumRep{} = False
isUnit SetRep{} = False
isUnit PredicateRep{} = False