-- | Apply some transformations to Michelson code.

module Michelson.Preprocess
  ( transformStrings
  , transformBytes
  ) where

import Data.Default (def)
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

import Michelson.Text (MText)
import Michelson.Typed

-- Note: we may add such transformation for long bytestrings as well if deemed necessary.
-- And for other constants which may be arbitrarily large (e. g. lists).
-- For now we need it only for strings and probably won't need for anything else.

-- | Transform all strings in a typed instructions using given
-- function. The first argument specifies whether we should go into
-- arguments that contain instructions.
transformStrings :: Bool -> (MText -> MText) -> Instr inp out -> Instr inp out
transformStrings goToValues f = transformConstants goToValues valF
  where
    valF :: Value t -> Value t
    valF = \case
      VString str -> VString $ f str
      VOption mv -> VOption (valF <$> mv)
      VList vs -> VList (valF <$> vs)
      VSet vSet -> VSet (Set.map valF vSet)
      VPair (v1, v2) -> VPair (valF v1, valF v2)
      VOr eith -> VOr (bimap valF valF eith)
      VMap m -> VMap . fmap valF . Map.mapKeys valF $ m
      VBigMap m -> VBigMap . fmap valF . Map.mapKeys valF $ m
      -- We do not handle 'VLam' here, because 'dfsInstr' takes care of that.
      v -> v

-- | Similar to 'transformStrings' but for bytes.
-- TODO [TM-375]: deduplicate
transformBytes :: Bool -> (ByteString -> ByteString) -> Instr inp out -> Instr inp out
transformBytes goToValues f = transformConstants goToValues valF
  where
    valF :: Value t -> Value t
    valF = \case
      VBytes str -> VBytes $ f str
      VOption mv -> VOption (valF <$> mv)
      VList vs -> VList (valF <$> vs)
      VSet vSet -> VSet (Set.map valF vSet)
      VPair (v1, v2) -> VPair (valF v1, valF v2)
      VOr eith -> VOr (bimap valF valF eith)
      VMap m -> VMap . fmap valF . Map.mapKeys valF $ m
      VBigMap m -> VBigMap . fmap valF . Map.mapKeys valF $ m
      -- We do not handle 'VLam' here, because 'dfsInstr' takes care of that.
      v -> v

transformConstants ::
     forall inp out.
     Bool
  -> (forall t. Value t -> Value t)
  -> Instr inp out
  -> Instr inp out
transformConstants dsGoToValues f = fst . dfsInstr def{ dsGoToValues } step
  where
    step :: forall i o. Instr i o -> (Instr i o, ())
    step = (,()) . \case
      PUSH v -> PUSH (f v)
      i -> i