{-# LANGUAGE LambdaCase #-}

-- | A module defining 'SigBindFamily' and other related types and
-- functions.
module HIndent.Pretty.SigBindFamily
  ( SigBindFamily(..)
  , LSigBindFamily
  , mkSortedLSigBindFamilyList
  , mkLSigBindFamilyList
  , filterLSig
  , filterLBind
  ) where

import Data.Function
import Data.List
import Data.Maybe
import GHC.Hs
import GHC.Types.SrcLoc

-- | A sum type containing one of those: function signature, function
-- binding, type family, type family instance, and data family instance.
data SigBindFamily
  = Sig (Sig GhcPs)
  | Bind (HsBindLR GhcPs GhcPs)
  | TypeFamily (FamilyDecl GhcPs)
  | TyFamInst (TyFamInstDecl GhcPs)
  | DataFamInst (DataFamInstDecl GhcPs)

-- | 'SigBindFamily' with the location information.
type LSigBindFamily = GenLocated SrcSpanAnnA SigBindFamily

-- | Creates a list of 'LSigBindFamily' from arguments. The list is sorted
-- by its elements' locations.
mkSortedLSigBindFamilyList ::
     [LSig GhcPs]
  -> [LHsBindLR GhcPs GhcPs]
  -> [LFamilyDecl GhcPs]
  -> [LTyFamInstDecl GhcPs]
  -> [LDataFamInstDecl GhcPs]
  -> [LSigBindFamily]
mkSortedLSigBindFamilyList :: [LSig GhcPs]
-> [LHsBindLR GhcPs GhcPs]
-> [LFamilyDecl GhcPs]
-> [LTyFamInstDecl GhcPs]
-> [LDataFamInstDecl GhcPs]
-> [LSigBindFamily]
mkSortedLSigBindFamilyList [LSig GhcPs]
sigs [LHsBindLR GhcPs GhcPs]
binds [LFamilyDecl GhcPs]
fams [LTyFamInstDecl GhcPs]
datafams =
  (LSigBindFamily -> LSigBindFamily -> Ordering)
-> [LSigBindFamily] -> [LSigBindFamily]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (RealSrcSpan -> RealSrcSpan -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (RealSrcSpan -> RealSrcSpan -> Ordering)
-> (LSigBindFamily -> RealSrcSpan)
-> LSigBindFamily
-> LSigBindFamily
-> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` SrcSpan -> RealSrcSpan
realSrcSpan (SrcSpan -> RealSrcSpan)
-> (LSigBindFamily -> SrcSpan) -> LSigBindFamily -> RealSrcSpan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SrcSpanAnnA -> SrcSpan
forall a. SrcSpanAnn' a -> SrcSpan
locA (SrcSpanAnnA -> SrcSpan)
-> (LSigBindFamily -> SrcSpanAnnA) -> LSigBindFamily -> SrcSpan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LSigBindFamily -> SrcSpanAnnA
forall l e. GenLocated l e -> l
getLoc) ([LSigBindFamily] -> [LSigBindFamily])
-> ([GenLocated SrcSpanAnnA (DataFamInstDecl GhcPs)]
    -> [LSigBindFamily])
-> [GenLocated SrcSpanAnnA (DataFamInstDecl GhcPs)]
-> [LSigBindFamily]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
  [LSig GhcPs]
-> [LHsBindLR GhcPs GhcPs]
-> [LFamilyDecl GhcPs]
-> [LTyFamInstDecl GhcPs]
-> [LDataFamInstDecl GhcPs]
-> [LSigBindFamily]
mkLSigBindFamilyList [LSig GhcPs]
sigs [LHsBindLR GhcPs GhcPs]
binds [LFamilyDecl GhcPs]
fams [LTyFamInstDecl GhcPs]
datafams

-- | Creates a list of 'LSigBindFamily' from arguments.
mkLSigBindFamilyList ::
     [LSig GhcPs]
  -> [LHsBindLR GhcPs GhcPs]
  -> [LFamilyDecl GhcPs]
  -> [LTyFamInstDecl GhcPs]
  -> [LDataFamInstDecl GhcPs]
  -> [LSigBindFamily]
mkLSigBindFamilyList :: [LSig GhcPs]
-> [LHsBindLR GhcPs GhcPs]
-> [LFamilyDecl GhcPs]
-> [LTyFamInstDecl GhcPs]
-> [LDataFamInstDecl GhcPs]
-> [LSigBindFamily]
mkLSigBindFamilyList [LSig GhcPs]
sigs [LHsBindLR GhcPs GhcPs]
binds [LFamilyDecl GhcPs]
fams [LTyFamInstDecl GhcPs]
insts [LDataFamInstDecl GhcPs]
datafams =
  (GenLocated SrcSpanAnnA (Sig GhcPs) -> LSigBindFamily)
-> [GenLocated SrcSpanAnnA (Sig GhcPs)] -> [LSigBindFamily]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Sig GhcPs -> SigBindFamily)
-> GenLocated SrcSpanAnnA (Sig GhcPs) -> LSigBindFamily
forall a b.
(a -> b) -> GenLocated SrcSpanAnnA a -> GenLocated SrcSpanAnnA b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Sig GhcPs -> SigBindFamily
Sig) [GenLocated SrcSpanAnnA (Sig GhcPs)]
[LSig GhcPs]
sigs [LSigBindFamily] -> [LSigBindFamily] -> [LSigBindFamily]
forall a. [a] -> [a] -> [a]
++
  (GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs) -> LSigBindFamily)
-> [GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs)]
-> [LSigBindFamily]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((HsBindLR GhcPs GhcPs -> SigBindFamily)
-> GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs) -> LSigBindFamily
forall a b.
(a -> b) -> GenLocated SrcSpanAnnA a -> GenLocated SrcSpanAnnA b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap HsBindLR GhcPs GhcPs -> SigBindFamily
Bind) [GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs)]
[LHsBindLR GhcPs GhcPs]
binds [LSigBindFamily] -> [LSigBindFamily] -> [LSigBindFamily]
forall a. [a] -> [a] -> [a]
++
  (GenLocated SrcSpanAnnA (FamilyDecl GhcPs) -> LSigBindFamily)
-> [GenLocated SrcSpanAnnA (FamilyDecl GhcPs)] -> [LSigBindFamily]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((FamilyDecl GhcPs -> SigBindFamily)
-> GenLocated SrcSpanAnnA (FamilyDecl GhcPs) -> LSigBindFamily
forall a b.
(a -> b) -> GenLocated SrcSpanAnnA a -> GenLocated SrcSpanAnnA b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap FamilyDecl GhcPs -> SigBindFamily
TypeFamily) [GenLocated SrcSpanAnnA (FamilyDecl GhcPs)]
[LFamilyDecl GhcPs]
fams [LSigBindFamily] -> [LSigBindFamily] -> [LSigBindFamily]
forall a. [a] -> [a] -> [a]
++
  (GenLocated SrcSpanAnnA (TyFamInstDecl GhcPs) -> LSigBindFamily)
-> [GenLocated SrcSpanAnnA (TyFamInstDecl GhcPs)]
-> [LSigBindFamily]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TyFamInstDecl GhcPs -> SigBindFamily)
-> GenLocated SrcSpanAnnA (TyFamInstDecl GhcPs) -> LSigBindFamily
forall a b.
(a -> b) -> GenLocated SrcSpanAnnA a -> GenLocated SrcSpanAnnA b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TyFamInstDecl GhcPs -> SigBindFamily
TyFamInst) [GenLocated SrcSpanAnnA (TyFamInstDecl GhcPs)]
[LTyFamInstDecl GhcPs]
insts [LSigBindFamily] -> [LSigBindFamily] -> [LSigBindFamily]
forall a. [a] -> [a] -> [a]
++ (GenLocated SrcSpanAnnA (DataFamInstDecl GhcPs) -> LSigBindFamily)
-> [GenLocated SrcSpanAnnA (DataFamInstDecl GhcPs)]
-> [LSigBindFamily]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((DataFamInstDecl GhcPs -> SigBindFamily)
-> GenLocated SrcSpanAnnA (DataFamInstDecl GhcPs) -> LSigBindFamily
forall a b.
(a -> b) -> GenLocated SrcSpanAnnA a -> GenLocated SrcSpanAnnA b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap DataFamInstDecl GhcPs -> SigBindFamily
DataFamInst) [GenLocated SrcSpanAnnA (DataFamInstDecl GhcPs)]
[LDataFamInstDecl GhcPs]
datafams

-- | Filters out 'Sig's and extract the wrapped values.
filterLSig :: [LSigBindFamily] -> [LSig GhcPs]
filterLSig :: [LSigBindFamily] -> [LSig GhcPs]
filterLSig =
  (LSigBindFamily -> Maybe (GenLocated SrcSpanAnnA (Sig GhcPs)))
-> [LSigBindFamily] -> [GenLocated SrcSpanAnnA (Sig GhcPs)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
    (\case
       (L SrcSpanAnnA
l (Sig Sig GhcPs
x)) -> GenLocated SrcSpanAnnA (Sig GhcPs)
-> Maybe (GenLocated SrcSpanAnnA (Sig GhcPs))
forall a. a -> Maybe a
Just (GenLocated SrcSpanAnnA (Sig GhcPs)
 -> Maybe (GenLocated SrcSpanAnnA (Sig GhcPs)))
-> GenLocated SrcSpanAnnA (Sig GhcPs)
-> Maybe (GenLocated SrcSpanAnnA (Sig GhcPs))
forall a b. (a -> b) -> a -> b
$ SrcSpanAnnA -> Sig GhcPs -> GenLocated SrcSpanAnnA (Sig GhcPs)
forall l e. l -> e -> GenLocated l e
L SrcSpanAnnA
l Sig GhcPs
x
       LSigBindFamily
_ -> Maybe (GenLocated SrcSpanAnnA (Sig GhcPs))
forall a. Maybe a
Nothing)

-- | Filters out 'Bind's and extract the wrapped values.
filterLBind :: [LSigBindFamily] -> [LHsBindLR GhcPs GhcPs]
filterLBind :: [LSigBindFamily] -> [LHsBindLR GhcPs GhcPs]
filterLBind =
  (LSigBindFamily
 -> Maybe (GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs)))
-> [LSigBindFamily]
-> [GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe
    (\case
       (L SrcSpanAnnA
l (Bind HsBindLR GhcPs GhcPs
x)) -> GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs)
-> Maybe (GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs))
forall a. a -> Maybe a
Just (GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs)
 -> Maybe (GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs)))
-> GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs)
-> Maybe (GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs))
forall a b. (a -> b) -> a -> b
$ SrcSpanAnnA
-> HsBindLR GhcPs GhcPs
-> GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs)
forall l e. l -> e -> GenLocated l e
L SrcSpanAnnA
l HsBindLR GhcPs GhcPs
x
       LSigBindFamily
_ -> Maybe (GenLocated SrcSpanAnnA (HsBindLR GhcPs GhcPs))
forall a. Maybe a
Nothing)