{-
  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.
-}

module Glean.Bytecode.Decode
  ( Decode
  , Decodable(..)
  , runDecode
  , decodeAll
  )
where

import Control.Monad
import Control.Monad.State.Strict (State, runState, state)
import Control.Monad.Trans.Maybe (MaybeT(..))
import Data.Word (Word64)

import Glean.Bytecode.Types

-- | A decoder which can be applied to a stream of words via 'runDecode'.
type Decode = MaybeT (State [Word64])

-- | Execute a 'Decode' and return the result and the remaining words.
runDecode :: Decode a -> [Word64] -> Maybe (a, [Word64])
runDecode (MaybeT f) xs = case runState f xs of
  (Just x, ys) -> Just (x, ys)
  _ -> Nothing

-- | Things that can be decoded.
class Decodable arg where
  decode :: Decode arg

instance Decodable Word64 where
  decode = MaybeT $ state $ \xs -> case xs of
    y : ys -> (Just y, ys)
    _ -> (Nothing, xs)

instance Decodable Literal where
  decode = Literal <$> decode

instance Decodable Label where
  decode = Label . fromIntegral <$> (decode :: Decode Word64)

instance Decodable [Label] where
  decode = do
    n <- decode
    replicateM (fromIntegral (n :: Word64)) decode

instance Decodable (Register t) where
  decode = Register <$> decode

-- | Decode as many elements as possible and return the decoded elements and
-- the rest of the words, starting with the first that can't be decoded.
decodeAll :: Decodable a => [Word64] -> ([a], [Word64])
decodeAll = go []
  where
    go xs [] = (reverse xs, [])
    go xs vs = case runDecode decode vs of
      Just (x, ws) -> go (x:xs) ws
      Nothing -> (reverse xs, vs)