{-# 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
  { CallGraph -> [Module]
_modules :: [Module],
    CallGraph -> Set (Key, Key)
_calls :: Set (Key, Key),
    CallGraph -> Set (Key, Key)
_types :: Set (Key, Key)
  }

data Module = Module
  { Module -> String
moduleName :: String,
    Module -> String
modulePath :: FilePath,
    Module -> Forest Decl
moduleForest :: Forest Decl
  }

data Decl = Decl
  { Decl -> String
declName :: String,
    Decl -> Key
declKey :: Key,
    Decl -> EnumSet GHCKey
declGHCKeys :: EnumSet GHCKey,
    Decl -> Bool
declExported :: Bool,
    Decl -> DeclType
declType :: DeclType,
    Decl -> Loc
declLoc :: Loc
  }

-- | A key in our own local space, c.f. a key that was generated by GHC.
newtype Key = Key {Key -> Int
unKey :: Int}
  deriving (Int -> Key
Key -> Int
Key -> [Key]
Key -> Key
Key -> Key -> [Key]
Key -> Key -> Key -> [Key]
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: Key -> Key -> Key -> [Key]
$cenumFromThenTo :: Key -> Key -> Key -> [Key]
enumFromTo :: Key -> Key -> [Key]
$cenumFromTo :: Key -> Key -> [Key]
enumFromThen :: Key -> Key -> [Key]
$cenumFromThen :: Key -> Key -> [Key]
enumFrom :: Key -> [Key]
$cenumFrom :: Key -> [Key]
fromEnum :: Key -> Int
$cfromEnum :: Key -> Int
toEnum :: Int -> Key
$ctoEnum :: Int -> Key
pred :: Key -> Key
$cpred :: Key -> Key
succ :: Key -> Key
$csucc :: Key -> Key
Enum, Int -> Key -> ShowS
[Key] -> ShowS
Key -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Key] -> ShowS
$cshowList :: [Key] -> ShowS
show :: Key -> String
$cshow :: Key -> String
showsPrec :: Int -> Key -> ShowS
$cshowsPrec :: Int -> Key -> ShowS
Show, Key -> Key -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Key -> Key -> Bool
$c/= :: Key -> Key -> Bool
== :: Key -> Key -> Bool
$c== :: Key -> Key -> Bool
Eq, Eq Key
Key -> Key -> Bool
Key -> Key -> Ordering
Key -> Key -> Key
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Key -> Key -> Key
$cmin :: Key -> Key -> Key
max :: Key -> Key -> Key
$cmax :: Key -> Key -> Key
>= :: Key -> Key -> Bool
$c>= :: Key -> Key -> Bool
> :: Key -> Key -> Bool
$c> :: Key -> Key -> Bool
<= :: Key -> Key -> Bool
$c<= :: Key -> Key -> Bool
< :: Key -> Key -> Bool
$c< :: Key -> Key -> Bool
compare :: Key -> Key -> Ordering
$ccompare :: Key -> Key -> Ordering
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 {GHCKey -> Int
unGHCKey :: Int}
  deriving newtype (Int -> GHCKey -> ShowS
[GHCKey] -> ShowS
GHCKey -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GHCKey] -> ShowS
$cshowList :: [GHCKey] -> ShowS
show :: GHCKey -> String
$cshow :: GHCKey -> String
showsPrec :: Int -> GHCKey -> ShowS
$cshowsPrec :: Int -> GHCKey -> ShowS
Show, Int -> GHCKey
GHCKey -> Int
GHCKey -> [GHCKey]
GHCKey -> GHCKey
GHCKey -> GHCKey -> [GHCKey]
GHCKey -> GHCKey -> GHCKey -> [GHCKey]
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: GHCKey -> GHCKey -> GHCKey -> [GHCKey]
$cenumFromThenTo :: GHCKey -> GHCKey -> GHCKey -> [GHCKey]
enumFromTo :: GHCKey -> GHCKey -> [GHCKey]
$cenumFromTo :: GHCKey -> GHCKey -> [GHCKey]
enumFromThen :: GHCKey -> GHCKey -> [GHCKey]
$cenumFromThen :: GHCKey -> GHCKey -> [GHCKey]
enumFrom :: GHCKey -> [GHCKey]
$cenumFrom :: GHCKey -> [GHCKey]
fromEnum :: GHCKey -> Int
$cfromEnum :: GHCKey -> Int
toEnum :: Int -> GHCKey
$ctoEnum :: Int -> GHCKey
pred :: GHCKey -> GHCKey
$cpred :: GHCKey -> GHCKey
succ :: GHCKey -> GHCKey
$csucc :: GHCKey -> GHCKey
Enum, GHCKey -> GHCKey -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: GHCKey -> GHCKey -> Bool
$c/= :: GHCKey -> GHCKey -> Bool
== :: GHCKey -> GHCKey -> Bool
$c== :: GHCKey -> GHCKey -> Bool
Eq, Eq GHCKey
GHCKey -> GHCKey -> Bool
GHCKey -> GHCKey -> Ordering
GHCKey -> GHCKey -> GHCKey
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: GHCKey -> GHCKey -> GHCKey
$cmin :: GHCKey -> GHCKey -> GHCKey
max :: GHCKey -> GHCKey -> GHCKey
$cmax :: GHCKey -> GHCKey -> GHCKey
>= :: GHCKey -> GHCKey -> Bool
$c>= :: GHCKey -> GHCKey -> Bool
> :: GHCKey -> GHCKey -> Bool
$c> :: GHCKey -> GHCKey -> Bool
<= :: GHCKey -> GHCKey -> Bool
$c<= :: GHCKey -> GHCKey -> Bool
< :: GHCKey -> GHCKey -> Bool
$c< :: GHCKey -> GHCKey -> Bool
compare :: GHCKey -> GHCKey -> Ordering
$ccompare :: GHCKey -> GHCKey -> Ordering
Ord)

data DeclType
  = ValueDecl
  | RecDecl
  | ConDecl
  | DataDecl
  | ClassDecl
  deriving
    (DeclType -> DeclType -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DeclType -> DeclType -> Bool
$c/= :: DeclType -> DeclType -> Bool
== :: DeclType -> DeclType -> Bool
$c== :: DeclType -> DeclType -> Bool
Eq, Eq DeclType
DeclType -> DeclType -> Bool
DeclType -> DeclType -> Ordering
DeclType -> DeclType -> DeclType
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: DeclType -> DeclType -> DeclType
$cmin :: DeclType -> DeclType -> DeclType
max :: DeclType -> DeclType -> DeclType
$cmax :: DeclType -> DeclType -> DeclType
>= :: DeclType -> DeclType -> Bool
$c>= :: DeclType -> DeclType -> Bool
> :: DeclType -> DeclType -> Bool
$c> :: DeclType -> DeclType -> Bool
<= :: DeclType -> DeclType -> Bool
$c<= :: DeclType -> DeclType -> Bool
< :: DeclType -> DeclType -> Bool
$c< :: DeclType -> DeclType -> Bool
compare :: DeclType -> DeclType -> Ordering
$ccompare :: DeclType -> DeclType -> Ordering
Ord, Int -> DeclType -> ShowS
[DeclType] -> ShowS
DeclType -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DeclType] -> ShowS
$cshowList :: [DeclType] -> ShowS
show :: DeclType -> String
$cshow :: DeclType -> String
showsPrec :: Int -> DeclType -> ShowS
$cshowsPrec :: Int -> DeclType -> ShowS
Show)

data Loc = Loc
  { Loc -> Int
locLine :: !Int,
    Loc -> Int
locCol :: !Int
  }
  deriving (Loc -> Loc -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Loc -> Loc -> Bool
$c/= :: Loc -> Loc -> Bool
== :: Loc -> Loc -> Bool
$c== :: Loc -> Loc -> Bool
Eq, Eq Loc
Loc -> Loc -> Bool
Loc -> Loc -> Ordering
Loc -> Loc -> Loc
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Loc -> Loc -> Loc
$cmin :: Loc -> Loc -> Loc
max :: Loc -> Loc -> Loc
$cmax :: Loc -> Loc -> Loc
>= :: Loc -> Loc -> Bool
$c>= :: Loc -> Loc -> Bool
> :: Loc -> Loc -> Bool
$c> :: Loc -> Loc -> Bool
<= :: Loc -> Loc -> Bool
$c<= :: Loc -> Loc -> Bool
< :: Loc -> Loc -> Bool
$c< :: Loc -> Loc -> Bool
compare :: Loc -> Loc -> Ordering
$ccompare :: Loc -> Loc -> Ordering
Ord)

instance Show Loc where
  showsPrec :: Int -> Loc -> ShowS
showsPrec Int
_ (Loc Int
ln Int
col) = forall a. Show a => a -> ShowS
shows Int
ln forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar Char
':' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows Int
col

{-# INLINE modDecls #-}
modDecls :: Traversal' Module Decl
modDecls :: Traversal' Module Decl
modDecls = Traversal' Module (Forest Decl)
modForest forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Traversal (Forest a) (Forest b) a b
forestT

{-# INLINE modForest #-}
modForest :: Traversal' Module (Forest Decl)
modForest :: Traversal' Module (Forest Decl)
modForest Forest Decl -> m (Forest Decl)
f (Module String
nm String
fp Forest Decl
ds) = String -> String -> Forest Decl -> Module
Module String
nm String
fp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Forest Decl -> m (Forest Decl)
f Forest Decl
ds

{-# INLINE forestT #-}
forestT :: Traversal (Forest a) (Forest b) a b
forestT :: forall a b. Traversal (Forest a) (Forest b) a b
forestT = forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse

rekeyCalls :: (Enum a, Ord b) => EnumMap a b -> Set (a, a) -> Set (b, b)
rekeyCalls :: forall a b.
(Enum a, Ord b) =>
EnumMap a b -> Set (a, a) -> Set (b, b)
rekeyCalls EnumMap a b
m = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. a -> a
id forall a. Ord a => a -> Set a -> Set a
Set.insert forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall k a. Enum k => k -> EnumMap k a -> Maybe a
EnumMap.lookup EnumMap a b
m) (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall k a. Enum k => k -> EnumMap k a -> Maybe a
EnumMap.lookup EnumMap a b
m)) forall a. Monoid a => a
mempty

ppCallGraph :: Prints CallGraph
ppCallGraph :: Prints CallGraph
ppCallGraph (CallGraph [Module]
modules Set (Key, Key)
_ Set (Key, Key)
_) = forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Module]
modules forall a b. (a -> b) -> a -> b
$ \(Module String
modName String
modPath Forest Decl
forest) -> do
  forall (m :: * -> *). MonadPrint m => String -> m ()
strLn forall a b. (a -> b) -> a -> b
$ String
modName forall a. Semigroup a => a -> a -> a
<> String
" (" forall a. Semigroup a => a -> a -> a
<> String
modPath forall a. Semigroup a => a -> a -> a
<> String
")"
  forall (m :: * -> *) a. MonadPrint m => m a -> m a
indent forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Prints (Tree Decl)
ppTree Forest Decl
forest

ppTree :: Prints (Tree Decl)
ppTree :: Prints (Tree Decl)
ppTree (Node (Decl String
name Key
_key EnumSet GHCKey
_ghckey Bool
_exp DeclType
typ Loc
loc) Forest Decl
children) = do
  forall (m :: * -> *). MonadPrint m => String -> m ()
strLn forall a b. (a -> b) -> a -> b
$ String
name forall a. Semigroup a => a -> a -> a
<> String
" (" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show DeclType
typ forall a. Semigroup a => a -> a -> a
<> String
", " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Loc
loc forall a. Semigroup a => a -> a -> a
<> String
")"
  forall (m :: * -> *) a. MonadPrint m => m a -> m a
indent forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Prints (Tree Decl)
ppTree Forest Decl
children