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.
  -}
  
  {-# LANGUAGE OverloadedStrings #-}
  -- | Generate "Glean.Schema"
  module Glean.Schema.Gen.Haskell
    ( genSchemaHS
    ) where
  
  import Data.HashMap.Strict (HashMap)
  import qualified Data.HashMap.Strict as HashMap
  import Data.Maybe
  import Data.Text (Text)
  import qualified Data.Text as Text
  import GHC.Stack (HasCallStack)
  import System.FilePath
  import TextShow
  
  import Glean.Schema.Gen.Utils
  import Glean.Angle.Types
  import Glean.Schema.Types
  
  genSchemaHS
    :: Version
    -> [ResolvedPredicateDef]
    -> [ResolvedTypeDef]
    -> Maybe Oncall
    -> [(FilePath,Text)]
  genSchemaHS _version preddefs typedefs oncall =
    ("hs" </> "TARGETS", genTargets declsPerNamespace oncall) :
    [ ("thrift" </> Text.unpack (underscored namespaces) <> "_include" <.> "hs",
        Text.intercalate (newline <> newline)
          (header namespaces deps : doGen preds types))
    | (namespaces, (deps, preds, types)) <- schemas
    ] ++
    [ ("hs" </> "Glean" </> "Schema" </>
        Text.unpack (Text.concat (map cap1 namespaces)) <.> "hs",
      genAllPredicates namespaces preds)
    | (namespaces, (_deps, preds, _types)) <- schemas
    ]
    where
      schemas = HashMap.toList declsPerNamespace
      declsPerNamespace =
        addNamespaceDependencies $ sortDeclsByNamespace preddefs typedefs
  
      namePolicy = mkNamePolicy preddefs typedefs
  
      doGen
        :: [ResolvedPredicateDef]
        -> [ResolvedTypeDef]
        -> [Text]
      doGen preds types = concat gen ++ reverse extra
        where
        (gen :: [[Text]], extra :: [Text]) = runM [] namePolicy typedefs $ do
           ps <- mapM genPredicate preds
           ts <- mapM (\TypeDef{..} -> genType typeDefRef typeDefType) types
           return (ps ++ ts)
  
  genTargets
    :: HashMap NameSpaces ([NameSpaces], [ResolvedPredicateDef], [ResolvedTypeDef])
    -> Maybe Oncall
    -> Text
  genTargets info oncall =
    Text.unlines $
       [ "# \x40generated"
       , "# to regenerate: ./glean/schema/sync"
       , "load(\"@fbcode_macros//build_defs:haskell_library.bzl\", " <>
         "\"haskell_library\")"
       , buckOncallAnnotation oncall
       , "" ] ++
       concatMap genTarget (HashMap.keys info)
    where
    genTarget ns =
      let
        namespace = underscored ns
      in
      -- mini Haskell library for the module containing allPredicates
      [ "haskell_library("
      , "  name = \"" <> namespace <> "\","
      , "  srcs = [\"Glean/Schema/" <> Text.concat (map cap1 ns) <>
          ".hs\"],"
      , "  deps = [\"//glean/if:glean-hs2\"]"
      , ")"
      , ""
      ]
  
  
  genAllPredicates
    :: NameSpaces
    -> [ResolvedPredicateDef]
    -> Text
  genAllPredicates namespace preds = Text.unlines $
    [ "-- @" <> "generated"
    , "module Glean.Schema." <>
      Text.concat (map cap1 namespace) <> " (allPredicates) where"
    , ""
    , "import Glean.Types"
    , ""
    , "allPredicates :: [PredicateRef]"
    , "allPredicates ="
    ] ++
    indentLines (encsep "[ " ", " "]"
      [ "(PredicateRef \"" <> predicateRef_name ref <> "\" " <>
           showt (predicateRef_version ref) <> ")"
      | pred <- preds
      , let ref = predicateDefRef pred
      ])
    where
      encsep start _ end [] = [start <> end]
      encsep start mid end xs =
        zipWith (<>) (start : repeat mid) xs ++ [end]
  
  header :: NameSpaces -> [NameSpaces] -> Text
  header here deps = Text.unlines $
    [ "-- @" <> "generated"
    , "{-# LANGUAGE TypeFamilies, MultiParamTypeClasses, DataKinds #-}"
    , "{-# LANGUAGE FlexibleInstances, TypeSynonymInstances #-}"
    , "{-# LANGUAGE UndecidableInstances #-}"
      -- UndecidableInstances is needed for the RecordFields type instances
    , "import qualified Data.ByteString"
    , "import qualified Data.Default"
    , "import qualified Data.Text"
    , ""
      -- we should use qualified imports as far as possible to avoid
      -- clashing with Thrift-generated code
    , "import qualified Glean.Types as Glean"
    , "import qualified Glean.Typed as Glean"
    , "import qualified Glean.Query.Angle as Angle"
    , "import qualified Glean.Angle.Types as Angle"
    , ""
    ] ++
    -- import dependencies
    map importSchema deps ++
    -- inject this instance into the Builtin schema, because there's no
    -- other good place to put it.
    (case here of
      ["builtin"] ->
        [ ""
        , "type instance Angle.SumFields (Prelude.Maybe t) ="
        , "  'Angle.TField \"nothing\" Unit ("
        , "  'Angle.TField \"just\" t"
        , "  'Angle.TNoFields)"
        ]
      _ -> [])
  
  
  importSchema :: NameSpaces -> Text
  importSchema ns =
    "import qualified Glean.Schema." <> upperSquashNS <> ".Types"
    where
    upperSquashNS = Text.concat (map cap1 ns)
  
  
  -- Check against hardcoded list of what glean.h provides
  provided :: Text -> Bool
  provided = (`elem` known)
    where known = []
  
  -- -----------------------------------------------------------------------------
  
  indent :: Text -> Text
  indent t = if Text.null t then t else "  " <> t
  
  indentLines :: [Text] -> [Text]
  indentLines = map indent
  
  -- | These names are for predicates/struct/union of Maybe-tree
  haskellTypeName :: (NameSpaces, Text) -> Text
  haskellTypeName (ns, x) = haskellThriftName (ns, cap1 x)
  
  fieldName :: (NameSpaces, Text) -> Text
  fieldName (ns, x) = haskellThriftName (ns, low1 x)
  
  -- | apply stupid and wrong heuristic
  paren :: Text -> Text
  paren inner =
    let needsParen = Text.any (' ' ==) inner
                     && (Text.take 1 inner `notElem` [ "(", "[" ])
    in if needsParen then "(" <> inner <> ")"
                     else inner
  
  -- | Slightly pretty breaking of long lines at "="
  (=@) :: Text -> Text -> Text
  (=@) lhs rhs = if Text.length lhs + Text.length rhs + 3 + 4 <= 80
                 then lhs <> " = " <> rhs
                 else let pre = Text.takeWhile (' '==) lhs
                      in lhs <> " =\n  " <> pre <> indent rhs
  infixr 5 =@
  
  optionalize :: Text -> Text
  optionalize name =
    let inner = fromMaybe name (Text.stripPrefix "inner " name)
    in "Prelude.Maybe" <>! inner
  
  die :: HasCallStack => String -> String -> a
  die x y = error ("Error in genSchemaHS in " <> x <> " : " <> y)
  
  (<>!) :: Text -> Text -> Text
  (<>!) a b = a <> " " <> paren b
  
  infixr 6 <>!
  
  genericParamNames :: Text -> Int -> [Text]
  genericParamNames base n = map mk [1..n]
    where mk i = base <> Text.pack (show i)
  
  localOrExternal :: NameSpaces -> Text -> (NameSpaces, Text)
  localOrExternal here name = if null ns then (here,x) else (ns,x)
    where (ns,x) = splitDot name
  
  shareTypeDef :: Bool -> NameSpaces -> ResolvedType' s -> M Text
  shareTypeDef genSub here t = do
    (no, name) <- nameThisType t
    case no of
      New | genSub -> do
        let tref = TypeRef (joinDot (here,name)) 0
        pushDefs =<< genType tref t
      _otherwise -> return ()
    return (haskellTypeName (localOrExternal here name))
  
  haskellTy :: NameSpaces -> ResolvedType' s -> M Text
  haskellTy = haskellTy_ PredName True
  
  -- | how to render predicate types in haskellTy
  data PredTy = PredName | PredKey
  
  haskellTy_
    :: PredTy
    -> Bool -- ^ generate nested type definitions
    -> NameSpaces
    -> ResolvedType' s
    -> M Text
  haskellTy_ withId genSub here t = case t of
    -- Leafs
    ByteTy{} -> return "Glean.Byte"
    NatTy{} -> return "Glean.Nat"
    BooleanTy{} -> return "Prelude.Bool"
    StringTy{} -> return "Data.Text.Text"
    ArrayTy ByteTy -> return "Data.ByteString.ByteString"
    ArrayTy tInner -> do
      inner <- haskellTy_ PredName genSub here tInner
      return $ "[" <> inner <> "]"
    RecordTy{} -> shareTypeDef genSub here t
    SumTy{} -> shareTypeDef genSub here t
    SetTy tInner -> do
      inner <- haskellTy_ PredName genSub here tInner
      return $ "[" <> inner <> "]"
    MaybeTy ty -> do
      inner <- haskellTy_ PredName genSub here ty
      return (optionalize inner)
    -- References
    PredicateTy _ pred -> do
      let wrap = case withId of
            PredName -> id
            PredKey -> ("Glean.KeyType " <>)
      wrap . haskellTypeName <$> predicateName pred
  
    NamedTy _ typeRef ->
      haskellTypeName <$> typeName typeRef
    EnumeratedTy _ -> shareTypeDef genSub here t
    TyVar{} -> error "haskellTy_: TyVar"
    HasTy{} -> error "haskellTy_: HasTy"
    HasKey{} -> error "haskellTy_: HasKey"
    ElementsOf{} -> error "haskellTy_: ElementsOf"
  
  
  genPredicate :: ResolvedPredicateDef -> M [Text]
  genPredicate PredicateDef{..}
    | provided (predicateRef_name predicateDefRef) = return []
    | otherwise = do
      pName <- predicateName predicateDefRef
      let
        here = fst pName
        name = haskellTypeName pName -- e.g. Clang_File
  
      withPredicateDefHint (snd pName) $ do
      let
        appendName suffix = let (ns, x) = pName in (ns, x <> suffix)
        glean_name = predicateRef_name predicateDefRef -- e.g. clang.File
        has_value = predicateDefValueType /= unitT
        field_id = fieldName $ appendName "_id" -- field name for Id
        name_key = appendName "_key" -- type of Key
        name_value = appendName "_value" -- type of Value
  
      (type_key, define_key) <-
        if shouldNameKeyType predicateDefKeyType then
          define_kt here predicateDefKeyType name_key
        else do
          ty <- haskellTy_ PredName True here predicateDefKeyType
          return (ty, [])
  
      (type_value, define_value) <-
        if not has_value then return ("Unit", []) else do
          define_kt here predicateDefValueType name_value
  
      let extra = define_key ++ define_value
      let ver = predicateRef_version predicateDefRef
          inst cls body =
            "instance " <> cls <> " " <> name <> " where"
            : indentLines body
          def_Predicate = inst "Glean.Predicate" $
            ["type KeyType " <> name =@ type_key]
            ++
            ["type ValueType " <> name =@ type_value | has_value]
            ++
            [ "getName _proxy " =@ "Glean.PredicateRef " <>
                  Text.pack (show glean_name) <> -- adds quotes, does escaping
                  showt ver
            , "getId = Glean.IdOf . Glean.Fid . " <> field_id
            , "mkFact (Glean.IdOf (Glean.Fid x)) k "
                <> (if has_value then "v" else "_")
                <> " = " <> name <> " x k"
                <> (if has_value then " v" else "")
            , "getFactKey = " <> fieldName name_key
            , if has_value
                then "getFactValue = " <> fieldName name_value
                else "getFactValue _ = Prelude.Just ()"
            ]
  
          def_Type = inst "Glean.Type"
            [ "buildRtsValue b = Glean.buildRtsValue b . Glean.getId"
            , "decodeRtsValue = Glean.decodeRef"
            , "decodeAsFact = Glean.decodeFact"
            , "sourceType = Glean.predicateSourceType"
            ]
  
      return $ extra ++ map myUnlines [def_Predicate, def_Type]
  
  
  -- Make the thriftTy type text, and the needed [Text] blocks
  define_kt
    :: HasCallStack
    => NameSpaces
    -> ResolvedType
    -> (NameSpaces, Text)
    -> M (Text, [Text])
  define_kt here typ name_kt = case typ of
    ByteTy{} -> leaf
    NatTy{} -> leaf
    StringTy{} -> leaf
    RecordTy [] -> leaf
    RecordTy _fields -> alias typ
    ArrayTy{} -> alias typ
    SetTy{} -> alias typ
    SumTy [] -> leaf
    SumTy _fields -> alias typ
    MaybeTy{} -> alias typ
    PredicateTy{} -> leaf
    NamedTy{} -> alias typ
    BooleanTy{} -> leaf
    _other -> die "define_kt" (show typ)
   where
     gname = joinDot name_kt
  
     leaf = (,) <$> return (haskellTypeName name_kt) <*> return []
  
     alias t = do
      ref <- haskellTy here (NamedTy () (TypeRef gname 0))
      def <- genType (TypeRef gname 0) t
      return (ref,def)
  
  genType :: TypeRef -> ResolvedType' s -> M [Text]
  genType TypeRef{..} ty
    | provided typeRef_name = return []
    | otherwise =
    case ty of
      RecordTy fields -> structDef typeRef_name typeRef_version fields
      SumTy fields -> unionDef typeRef_name typeRef_version fields
      EnumeratedTy vals -> enumDef typeRef_name typeRef_version vals
      _ -> return []
  
  structDef :: Name -> Version -> [ResolvedFieldDef' s] -> M [Text]
  structDef ident ver fields = do
    let typeRef = TypeRef ident ver
    sName@(here,root) <- typeName typeRef
    let name = haskellTypeName sName
  
    withTypeDefHint root $ do
    let
      fieldParamNames = genericParamNames "x" (length fields)
      makeTypeName (FieldDef p tField) =
        withRecordFieldHint p (haskellTy_ PredKey True here tField)
      spaced = Text.intercalate " " fieldParamNames
      nameAndParams = if null fieldParamNames then name
                      else name <> " " <> spaced
      encodeMe = case length fields of
        0 -> [ "buildRtsValue _b " <> name <> " = Prelude.return ()" ]
        _ -> ("buildRtsValue b " <> paren nameAndParams <> " = do")
              : indentLines (map ("Glean.buildRtsValue b " <>) fieldParamNames)
      decodeMe = case length fields of
        0 -> [ "decodeRtsValue = Prelude.pure " <> name]
        n -> "decodeRtsValue = " <> name
             : indent "<$> Glean.decodeRtsValue"
             : indentLines (replicate (pred n) "<*> Glean.decodeRtsValue")
  
    tys <- mapM makeTypeName fields
    let def_Type =
          "instance Glean.Type " <> name <> " where"
          : indentLines
          ( encodeMe
          <> decodeMe
          <> [sourceTypeDef ident ver]
          )
  
        def_RecordFields =
          [ emitFieldTypes "Angle.RecordFields" name (zip fields tys) ]
  
    return $ map myUnlines [def_Type, def_RecordFields]
  
  
  unionDef :: Name -> Version -> [ResolvedFieldDef' s] -> M [Text]
  unionDef ident ver fields = do
    let typeRef = TypeRef ident ver
    uName@(here,root) <- typeName typeRef
    let name = haskellTypeName uName
    withTypeDefHint root $ do
    let
      shortConNames = map fieldDefName fields
      conNames = map toConName shortConNames
      toConName shortName = cap1 (prefix <> shortName)
        where prefix = name <> "_"
      makeTypeName pred gen (FieldDef p tField) =
        withUnionFieldHint p (haskellTy_ pred gen here tField)
  
    -- walk over the fields to generate nested types first
    mapM_ (makeTypeName PredName True) fields
  
    keyTypeNames <- mapM (makeTypeName PredKey False) fields
  
    let def_Type = case conNames of
          [] ->
            "instance Glean.Type " <> name <> " where" : indentLines
              [ "buildRtsValue _ _ = Prelude.return ()"
              , "decodeRtsValue = Prelude.error $ \"decodeRtsValue\" <> " <> name
              ]
          _ ->
            let emptyCon = toConName "EMPTY"
                builds =
                    mkBuildEmpty conNames emptyCon ++
                    concat (zipWith mkBuild [0..] conNames)
                decodes =
                  "decodeRtsValue = Glean.sumD" :
                  indentLines
                    ( "(Prelude.pure " <> emptyCon <> ")"
                    : asArray (map mkDecode conNames)
                    )
            in concat
              [ ["instance Glean.Type " <> name <> " where"]
              , indentLines builds
              , indentLines decodes
              , indentLines [sourceTypeDef ident ver]
              ]
  
        def_SumFields =
          [ emitFieldTypes "Angle.SumFields" name (zip fields keyTypeNames) ]
  
    return $ map myUnlines $ [def_Type, def_SumFields]
    where
      mkBuildEmpty constructors emptyCon =
        let index = length constructors in
        "buildRtsValue b " <> emptyCon <> "=" : indentLines
          [ "Glean.buildRtsSelector b " <> Text.pack (show index) ]
  
      mkBuild :: Int -> Text -> [Text]
      mkBuild i c =
        "buildRtsValue b (" <> c <> " x) = do" : indentLines
          [ "Glean.buildRtsSelector b " <> Text.pack (show i)
          , "Glean.buildRtsValue b x" ]
  
      mkDecode :: Text -> Text
      mkDecode c = "Glean.mapD " <> c
  
      asArray :: [Text] -> [Text]
      asArray = \case
        [] -> ["[]"]
        (x:xs) -> concat
          [ ["[ " <> x]
          , [", " <> e | e <- xs]
          , ["]"]
          ]
  
  enumDef :: Name -> Version -> [Name] -> M [Text]
  enumDef ident ver eVals = do
    let typeRef = TypeRef ident ver
    eName@(_,root) <- typeName typeRef
    let name = haskellTypeName eName
  
    withTypeDefHint root $ do
    let
      def_Type = "instance Glean.Type " <> name <> " where"
              : indentLines [ "buildRtsValue = Glean.thriftEnum_buildRtsValue "
                            , "decodeRtsValue = Glean.thriftEnumD "
                            , sourceTypeDef ident ver
                            ]
      def_SumFields =
         [ emitFieldTypes "Angle.SumFields" name
             [ (FieldDef n unitT, "Glean.Schema.Builtin.Types.Unit")
             | n <- eVals ]
         ]
  
      def_AngleEnum =
        "instance Angle.AngleEnum " <> name <> " where": indentLines
          [ "type AngleEnumTy " <> name <> " = " <> name
          , "enumName v = Text.pack (Prelude.drop " <>
              showt (Text.length root + 1) <> " (Prelude.show v))"]
  
    return $ map myUnlines [def_Type, def_SumFields, def_AngleEnum]
  
  sourceTypeDef :: Name -> Version -> Text
  sourceTypeDef name version =
    "sourceType _ = Angle.NamedTy () " <> paren sourceRef
    where
      sourceRef = Text.unwords
        [ "Angle.SourceRef"
        , Text.pack (show name)
        , paren ("Prelude.Just " <> showt version)
        ]
  
  
  emitFieldTypes :: Text -> Text -> [(ResolvedFieldDef' s, Text)] -> Text
  emitFieldTypes family name fields =
    "type instance " <> family <> " " <> name <> " = " <> go fields
    where
      go [] = "'Angle.TNoFields"
      go ((FieldDef name _, ty) : rest) =
         "'Angle.TField \"" <> name <> "\" (" <> ty <> ") (" <> go rest <> ")"
  