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.
  -}
  
  module Thrift.ExactPrint.Codemod
    ( roundTripWith
    , codemodConsts
    , lookupEnum
    ) where
  
  import qualified Data.Map.Strict as Map
  import Data.Some
  import qualified Data.Text.Lazy as Text
  import Data.Text (Text)
  
  import Thrift.Compiler.Options
  import Thrift.Compiler.Parser
  import Thrift.Compiler.Pretty
  import Thrift.Compiler.Typechecker
  import qualified Thrift.Compiler.Typechecker.Monad as Typechecker
  import Thrift.Compiler.Types
  import Thrift.Compiler.Plugins.Linter
  
  import Thrift.ExactPrint.Convert
  import Thrift.ExactPrint.PrettyPrint
  import Thrift.ExactPrint.Types
  
  -- | Round Trip a file, applying some tranformation to the typechecked AST
  roundTripWith
    :: (Program Linter Offset -> Program Linter Offset)
    -> String
    -> String
  roundTripWith f input = case runParser parseThrift "" input of
    Left e -> error e
    Right tf -> case typecheck (defaultOptions NoOpts) (mkModuleMap tf) of
      Left es -> error $ concatMap renderTypeError es
      Right (prog,_) -> Text.unpack $ exactPrint $ f $ computeOffsets prog
  
  mkModuleMap :: ([Parsed Header], [Parsed Decl]) -> ModuleMap
  mkModuleMap (headers, decls) =
    Map.singleton "" ThriftFile
      { thriftName = ""
      , thriftPath = ""
      , thriftHeaders = headers
      , thriftDecls   = decls
      , thriftSplice = Nothing
      , thriftComments = []
      }
  
  -- | Type aware codemod for constant values everywhere in the AST
  codemodConsts
    :: forall l a. (forall t. Env l -> Type l t -> ConstVal a -> ConstVal a)
    -> Program l a
    -> Program l a
  codemodConsts f p@Program{..} = p { progDecls = map codemodDecl progDecls }
    where
      codemodDecl (D_Struct s@Struct{..}) =
        D_Struct s { structMembers = map codemodField structMembers }
      codemodDecl (D_Const c@Const{..}) =
        D_Const c { constVal = codemodConst constResolvedType constVal }
      -- These have no constants
      codemodDecl d@D_Enum{} = d
      codemodDecl d@D_Typedef{} = d
      codemodDecl d@D_Union{} = d
      -- Services not supported yet
      codemodDecl d@D_Service{} = d
      -- Interactions not supported yet
      codemodDecl d@D_Interaction{} = d
  
      codemodField :: Field u 'Resolved l a -> Field u 'Resolved l a
      codemodField fl@Field{..} =
        fl { fieldVal = codemodConst fieldResolvedType <$> fieldVal }
  
      codemodConst :: Type l t -> UntypedConst a -> UntypedConst a
      codemodConst ty c@UntypedConst{..} =
        c { ucConst = codemodConstVal ty ucConst }
  
      -- Recursively apply the refactoring function to everything
      codemodConstVal :: Type l t -> ConstVal a -> ConstVal a
  
      -- List Types
      codemodConstVal t@(TList u) c = codemodListVal t u c
      codemodConstVal t@(TSet u) c = codemodListVal t u c
      codemodConstVal t@(THashSet u) c = codemodListVal t u c
  
      -- Map Types
      codemodConstVal t@(TMap k v) c = codemodMapVal t k v c
      codemodConstVal t@(THashMap k v) c = codemodMapVal t k v c
  
      -- Aliases
      codemodConstVal t@(TTypedef _ u _loc) c = f progEnv t $ codemodConstVal u c
      codemodConstVal t@(TNewtype _ u _loc) c = f progEnv t $ codemodConstVal u c
  
      -- Structs
      codemodConstVal t@(TStruct name _loc) c
        | Just (Some schema) <- lookupSchema name progEnv =
            f progEnv t $ codemodSchema schema c
      codemodConstVal t@(TException name _loc) c
        | Just (Some schema) <- lookupSchema name progEnv =
            f progEnv t $ codemodSchema schema c
      codemodConstVal t@(TUnion name _loc) c
        | Just (Some schema) <- lookupUnion name progEnv =
            f progEnv t $ codemodSchema schema c
  
      -- Everything Else
      codemodConstVal t c = f progEnv t c
  
      -- Apply the refactor inside the list and then outside the list
      codemodListVal :: Type l (f t) -> Type l t -> ConstVal a -> ConstVal a
      codemodListVal outer inner c@ListConst{..} = f progEnv outer
        c { lvElems = [ e { leElem = codemodConst inner leElem }
                      | e@ListElem{..} <- lvElems
                      ]
          }
      codemodListVal t _ c = f progEnv t c
  
      codemodMapVal
        :: Type l (f k v) -> Type l k -> Type l v -> ConstVal a -> ConstVal a
      codemodMapVal outer k v m@MapConst{..} = f progEnv outer
        m { mvElems =
             [ e { leElem = mp { mpKey = codemodConst k mpKey
                               , mpVal = codemodConst v mpVal
                               }
                 }
             | e@ListElem{leElem=mp@MapPair{..}} <- mvElems
             ]
          }
      codemodMapVal t _ _ c = f progEnv t c
  
      codemodSchema :: SCHEMA l t s -> ConstVal a -> ConstVal a
      codemodSchema schema m@MapConst{..} =
        m { mvElems =
               [ e { leElem = mp { mpVal =
                                      codemodSchemaField schema s mpVal
                                 }
                   }
               | e@ListElem
                   {leElem=
                       mp@MapPair{mpKey=UntypedConst{ucConst=StringConst s _},..}}
                   <- mvElems
               ]
          }
      codemodSchema _ c = c
  
      codemodSchemaField
        :: SCHEMA l t s -> Text -> UntypedConst a -> UntypedConst a
      codemodSchemaField SEmpty _ c = c
      codemodSchemaField (SField _ name ty schema) fl c
        | name == fl = codemodConst ty c
        | otherwise  = codemodSchemaField schema fl c
      codemodSchemaField (SReqField _ name ty schema) fl c
        | name == fl = codemodConst ty c
        | otherwise  = codemodSchemaField schema fl c
      codemodSchemaField (SOptField _ name ty schema) fl c
        | name == fl = codemodConst ty c
        | otherwise  = codemodSchemaField schema fl c
  
  lookupSchema :: Name -> Env l -> Maybe (Some (Schema l))
  lookupSchema = wrapLookup Typechecker.lookupSchema
  
  lookupUnion :: Name -> Env l -> Maybe (Some (USchema l))
  lookupUnion = wrapLookup Typechecker.lookupUnion
  
  lookupEnum :: Name -> Env l -> Maybe EnumValues
  lookupEnum = wrapLookup Typechecker.lookupEnum
  
  wrapLookup
    :: (ThriftName -> Loc -> Typechecker.TC l a)
    -> Name
    -> Env l
    -> Maybe a
  wrapLookup f Name{..} env =
    either (const Nothing) Just $ Typechecker.runTypechecker env $
    f sourceName noLoc
  