{-# 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.Util.Lens
import Calligraphy.Util.Printer
import Control.Monad
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 _ _) = forM_ modules $ \(Module modName modPath forest) -> do
  strLn $ modName <> " (" <> modPath <> ")"
  indent $ mapM_ ppTree forest

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