{-# 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]
(Key -> Key)
-> (Key -> Key)
-> (Int -> Key)
-> (Key -> Int)
-> (Key -> [Key])
-> (Key -> Key -> [Key])
-> (Key -> Key -> [Key])
-> (Key -> Key -> Key -> [Key])
-> Enum 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
(Int -> Key -> ShowS)
-> (Key -> String) -> ([Key] -> ShowS) -> Show Key
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
(Key -> Key -> Bool) -> (Key -> Key -> Bool) -> Eq Key
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
Eq Key
-> (Key -> Key -> Ordering)
-> (Key -> Key -> Bool)
-> (Key -> Key -> Bool)
-> (Key -> Key -> Bool)
-> (Key -> Key -> Bool)
-> (Key -> Key -> Key)
-> (Key -> Key -> Key)
-> Ord 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
$cp1Ord :: Eq Key
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
(Int -> GHCKey -> ShowS)
-> (GHCKey -> String) -> ([GHCKey] -> ShowS) -> Show GHCKey
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]
(GHCKey -> GHCKey)
-> (GHCKey -> GHCKey)
-> (Int -> GHCKey)
-> (GHCKey -> Int)
-> (GHCKey -> [GHCKey])
-> (GHCKey -> GHCKey -> [GHCKey])
-> (GHCKey -> GHCKey -> [GHCKey])
-> (GHCKey -> GHCKey -> GHCKey -> [GHCKey])
-> Enum 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
(GHCKey -> GHCKey -> Bool)
-> (GHCKey -> GHCKey -> Bool) -> Eq GHCKey
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
Eq GHCKey
-> (GHCKey -> GHCKey -> Ordering)
-> (GHCKey -> GHCKey -> Bool)
-> (GHCKey -> GHCKey -> Bool)
-> (GHCKey -> GHCKey -> Bool)
-> (GHCKey -> GHCKey -> Bool)
-> (GHCKey -> GHCKey -> GHCKey)
-> (GHCKey -> GHCKey -> GHCKey)
-> Ord 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
$cp1Ord :: Eq GHCKey
Ord)

data DeclType
  = ValueDecl
  | RecDecl
  | ConDecl
  | DataDecl
  | ClassDecl
  deriving
    (DeclType -> DeclType -> Bool
(DeclType -> DeclType -> Bool)
-> (DeclType -> DeclType -> Bool) -> Eq DeclType
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
Eq DeclType
-> (DeclType -> DeclType -> Ordering)
-> (DeclType -> DeclType -> Bool)
-> (DeclType -> DeclType -> Bool)
-> (DeclType -> DeclType -> Bool)
-> (DeclType -> DeclType -> Bool)
-> (DeclType -> DeclType -> DeclType)
-> (DeclType -> DeclType -> DeclType)
-> Ord 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
$cp1Ord :: Eq DeclType
Ord, Int -> DeclType -> ShowS
[DeclType] -> ShowS
DeclType -> String
(Int -> DeclType -> ShowS)
-> (DeclType -> String) -> ([DeclType] -> ShowS) -> Show DeclType
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
(Loc -> Loc -> Bool) -> (Loc -> Loc -> Bool) -> Eq Loc
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
Eq Loc
-> (Loc -> Loc -> Ordering)
-> (Loc -> Loc -> Bool)
-> (Loc -> Loc -> Bool)
-> (Loc -> Loc -> Bool)
-> (Loc -> Loc -> Bool)
-> (Loc -> Loc -> Loc)
-> (Loc -> Loc -> Loc)
-> Ord 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
$cp1Ord :: Eq Loc
Ord)

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

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

{-# INLINE modForest #-}
modForest :: Traversal' Module (Forest Decl)
modForest :: (Forest Decl -> m (Forest Decl)) -> Module -> m Module
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 (Forest Decl -> Module) -> m (Forest Decl) -> m Module
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 :: (a -> m b) -> Forest a -> m (Forest b)
forestT = (Tree a -> m (Tree b)) -> Forest a -> m (Forest b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((Tree a -> m (Tree b)) -> Forest a -> m (Forest b))
-> ((a -> m b) -> Tree a -> m (Tree b))
-> (a -> m b)
-> Forest a
-> m (Forest b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> m b) -> Tree a -> m (Tree b)
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 :: EnumMap a b -> Set (a, a) -> Set (b, b)
rekeyCalls EnumMap a b
m = ((a, a) -> Set (b, b) -> Set (b, b))
-> Set (b, b) -> Set (a, a) -> Set (b, b)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Set (b, b) -> Set (b, b))
-> ((b, b) -> Set (b, b) -> Set (b, b))
-> Maybe (b, b)
-> Set (b, b)
-> Set (b, b)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set (b, b) -> Set (b, b)
forall a. a -> a
id (b, b) -> Set (b, b) -> Set (b, b)
forall a. Ord a => a -> Set a -> Set a
Set.insert (Maybe (b, b) -> Set (b, b) -> Set (b, b))
-> ((a, a) -> Maybe (b, b)) -> (a, a) -> Set (b, b) -> Set (b, b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Maybe b) -> (a -> Maybe b) -> (a, a) -> Maybe (b, b)
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 ((a -> EnumMap a b -> Maybe b) -> EnumMap a b -> a -> Maybe b
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> EnumMap a b -> Maybe b
forall k a. Enum k => k -> EnumMap k a -> Maybe a
EnumMap.lookup EnumMap a b
m) ((a -> EnumMap a b -> Maybe b) -> EnumMap a b -> a -> Maybe b
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> EnumMap a b -> Maybe b
forall k a. Enum k => k -> EnumMap k a -> Maybe a
EnumMap.lookup EnumMap a b
m)) Set (b, b)
forall a. Monoid a => a
mempty

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

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