content: {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  
  module Calligraphy.Util.Types
    ( -- * Data types
      CallGraph (..),
      Module (..),
      Decl (..),
      DeclType (..),
      Key (..),
      GHCKey (..),
      Loc (..),
  
      -- * Utility functions
      rekeyCalls,
      ppCallGraph,
  
      -- * Lensy stuff
      over,
      forT_,
      modForest,
      modDecls,
      forestT,
    )
  where
  
  import Calligraphy.Prelude hiding (Decl, DeclType, Node)
  import Calligraphy.Util.Lens
  import Calligraphy.Util.Printer
  import Data.Bitraversable (bitraverse)
  import Data.EnumMap (EnumMap)
  import qualified Data.EnumMap as EnumMap
  import Data.EnumSet (EnumSet)
  import Data.Graph
  import Data.Set (Set)
  import qualified Data.Set as Set
  
  -- | This is the main type that processing phases will operate on.
  -- Note that calls and typing judgments are part of this top-level structure, not of the individual modules.
  data CallGraph = CallGraph
    { _modules :: [Module],
      _calls :: Set (Key, Key),
      _types :: Set (Key, Key)
    }
  
  data Module = Module
    { moduleName :: String,
      modulePath :: FilePath,
      moduleForest :: Forest Decl
    }
  
  data Decl = Decl
    { declName :: String,
      declKey :: Key,
      declGHCKeys :: EnumSet GHCKey,
      declExported :: Bool,
      declType :: DeclType,
      declLoc :: Loc
    }
  
  -- | A key in our own local space, c.f. a key that was generated by GHC.
  newtype Key = Key {unKey :: Int}
    deriving (Enum, Show, Eq, Ord)
  
  -- | A key that was produced by GHC, c.f. Key that we produced ourselves.
  -- We wrap it in a newtype because GHC itself uses a type synonym, but we want conversions to be as explicit as possible.
  newtype GHCKey = GHCKey {unGHCKey :: Int}
    deriving newtype (Show, Enum, Eq, Ord)
  
  data DeclType
    = ValueDecl
    | RecDecl
    | ConDecl
    | DataDecl
    | ClassDecl
    deriving
      (Eq, Ord, Show)
  
  data Loc = Loc
    { locLine :: !Int,
      locCol :: !Int
    }
    deriving (Eq, Ord)
  
  instance Show Loc where
    showsPrec _ (Loc ln col) = shows ln . showChar ':' . shows col
  
  {-# INLINE modDecls #-}
  modDecls :: Traversal' Module Decl
  modDecls = modForest . forestT
  
  {-# INLINE modForest #-}
  modForest :: Traversal' Module (Forest Decl)
  modForest f (Module nm fp ds) = Module nm fp <$> f ds
  
  {-# INLINE forestT #-}
  forestT :: Traversal (Forest a) (Forest b) a b
  forestT = traverse . traverse
  
  rekeyCalls :: (Enum a, Ord b) => EnumMap a b -> Set (a, a) -> Set (b, b)
  rekeyCalls m =
    foldr
      (maybe id Set.insert . bitraverse (flip EnumMap.lookup m) (flip EnumMap.lookup m))
      mempty
  
  ppCallGraph :: Prints CallGraph
  ppCallGraph (CallGraph modules calls types) = do
    forM_ modules $ \(Module modName modPath forest) -> do
      strLn $ modName <> " (" <> modPath <> ")"
      indent $ mapM_ ppTree forest
    let edgeArrows = Set.map showCall calls <> Set.map showType types
    forM_ edgeArrows strLn
      where
        table = makeDeclarationLookupTable (concatMap moduleForest modules)
        nameOfKey = declName . (EnumMap.!) table
        showCall (from, to) = nameOfKey from <> " ---> " <> nameOfKey to
        showType (from, to) = nameOfKey from <> " ···> " <> nameOfKey to
  
  ppTree :: Prints (Tree Decl)
  ppTree (Node (Decl name _key _ghckey _exp typ loc) children) = do
    strLn $ name <> " (" <> show typ <> ", " <> show loc <> ")"
    indent $ mapM_ ppTree children
  
  makeDeclarationLookupTable :: Forest Decl -> EnumMap Key Decl
  makeDeclarationLookupTable = (foldMap . foldMap) makeSingletonMap
    where
      makeSingletonMap declaration = EnumMap.singleton (declKey declaration) declaration
  