    {-# LANGUAGE OverloadedStrings #-}
    {-# LANGUAGE PatternSynonyms #-}
    {-# LANGUAGE RecordWildCards #-}
    {-# LANGUAGE TemplateHaskell #-}

    -- | Definitions for fixity analysis.
    module Ormolu.Fixity
      ( OpName,
        pattern OpName,
        unOpName,
        occOpName,
        FixityDirection (..),
        FixityInfo (..),
        defaultFixityInfo,
        FixityApproximation (..),
        defaultFixityApproximation,
        FixityOverrides (..),
        defaultFixityOverrides,
        ModuleReexports (..),
        defaultModuleReexports,
        PackageFixityMap (..),
        ModuleFixityMap (..),
        inferFixity,
        HackageInfo (..),
        hackageInfo,
        defaultDependencies,
        packageFixityMap,
        packageFixityMap',
        moduleFixityMap,
        applyFixityOverrides,
      )
    where

    import Data.Binary qualified as Binary
    import Data.Binary.Get qualified as Binary
    import Data.ByteString.Lazy qualified as BL
    import Data.FileEmbed (embedFile)
    import Data.List.NonEmpty (NonEmpty)
    import Data.List.NonEmpty qualified as NE
    import Data.Map.Strict qualified as Map
    import Data.MemoTrie (memo)
    import Data.Set (Set)
    import Data.Set qualified as Set
    import Distribution.ModuleName (ModuleName)
    import Distribution.Types.PackageName (PackageName, mkPackageName, unPackageName)
    import Language.Haskell.Syntax.ImpExp (ImportListInterpretation (..))
    import Ormolu.Fixity.Imports (FixityImport (..))
    import Ormolu.Fixity.Internal

    -- | The built-in 'HackageInfo' used by Ormolu.
    hackageInfo :: HackageInfo
    hackageInfo :: HackageInfo
    hackageInfo =
      Get HackageInfo -> ByteString -> HackageInfo
    forall a. Get a -> ByteString -> a
    Binary.runGet Get HackageInfo
    forall t. Binary t => Get t
    Binary.get (ByteString -> HackageInfo) -> ByteString -> HackageInfo
    forall a b. (a -> b) -> a -> b
    $
        StrictByteString -> ByteString
    BL.fromStrict $(embedFile "extract-hackage-info/hackage-info.bin")

    -- | Default set of packages to assume as dependencies e.g. when no Cabal
    -- file is found or taken into consideration.
    defaultDependencies :: Set PackageName
    defaultDependencies :: Set PackageName
    defaultDependencies = PackageName -> Set PackageName
    forall a. a -> Set a
    Set.singleton (String -> PackageName
    mkPackageName String
    "base")

    -- | Compute the fixity map that is specific to the package we are formatting.
    packageFixityMap ::
      -- | Set of packages to select
      Set PackageName ->
      -- | Package fixity map
      PackageFixityMap
    packageFixityMap :: Set PackageName -> PackageFixityMap
    packageFixityMap = HackageInfo -> Set PackageName -> PackageFixityMap
    packageFixityMap' HackageInfo
    hackageInfo

    -- | The same as 'packageFixityMap', except this specific version of the
    -- function allows the user to specify 'HackageInfo' used to build the final
    -- fixity map.
    packageFixityMap' ::
      -- | Hackage info
      HackageInfo ->
      -- | Set of packages to select
      Set PackageName ->
      -- | Package fixity map
      PackageFixityMap
    packageFixityMap' :: HackageInfo -> Set PackageName -> PackageFixityMap
    packageFixityMap' (HackageInfo Map PackageName (Map ModuleName (Map OpName FixityInfo))
    m) = (Set PackageName -> PackageFixityMap)
    -> Set PackageName -> PackageFixityMap
    forall v. (Set PackageName -> v) -> Set PackageName -> v
    memoSet ((Set PackageName -> PackageFixityMap)
     -> Set PackageName -> PackageFixityMap)
    -> (Set PackageName -> PackageFixityMap)
    -> Set PackageName
    -> PackageFixityMap
    forall a b. (a -> b) -> a -> b
    $ \Set PackageName
    dependencies ->
      -- The core idea here is to transform:
      --
      -- Map PackageName (Map ModuleName (Map OpName FixityInfo))
      --
      -- into
      --
      -- Map OpName [(PackageName, ModuleName, FixityInfo)]
      --
      -- which we accomplish by turning 'Map's into tuples with 'Map.toList' and
      -- then flattening them with 'flatten :: [(a, [b])] -> [(a, b)]'.
      --
      -- The target type results from the need to be able to quickly index by
      -- the operator name when we do fixity resolution later.
      Map OpName (NonEmpty (PackageName, ModuleName, FixityInfo))
    -> PackageFixityMap
    PackageFixityMap
        (Map OpName (NonEmpty (PackageName, ModuleName, FixityInfo))
     -> PackageFixityMap)
    -> (Map PackageName (Map ModuleName (Map OpName FixityInfo))
        -> Map OpName (NonEmpty (PackageName, ModuleName, FixityInfo)))
    -> Map PackageName (Map ModuleName (Map OpName FixityInfo))
    -> PackageFixityMap
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([(PackageName, ModuleName, FixityInfo)]
     -> Maybe (NonEmpty (PackageName, ModuleName, FixityInfo)))
    -> Map OpName [(PackageName, ModuleName, FixityInfo)]
    -> Map OpName (NonEmpty (PackageName, ModuleName, FixityInfo))
    forall a b k. (a -> Maybe b) -> Map k a -> Map k b
    Map.mapMaybe [(PackageName, ModuleName, FixityInfo)]
    -> Maybe (NonEmpty (PackageName, ModuleName, FixityInfo))
    forall a. [a] -> Maybe (NonEmpty a)
    NE.nonEmpty
        (Map OpName [(PackageName, ModuleName, FixityInfo)]
     -> Map OpName (NonEmpty (PackageName, ModuleName, FixityInfo)))
    -> (Map PackageName (Map ModuleName (Map OpName FixityInfo))
        -> Map OpName [(PackageName, ModuleName, FixityInfo)])
    -> Map PackageName (Map ModuleName (Map OpName FixityInfo))
    -> Map OpName (NonEmpty (PackageName, ModuleName, FixityInfo))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([(PackageName, ModuleName, FixityInfo)]
     -> [(PackageName, ModuleName, FixityInfo)]
     -> [(PackageName, ModuleName, FixityInfo)])
    -> [(OpName, [(PackageName, ModuleName, FixityInfo)])]
    -> Map OpName [(PackageName, ModuleName, FixityInfo)]
    forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
    Map.fromListWith [(PackageName, ModuleName, FixityInfo)]
    -> [(PackageName, ModuleName, FixityInfo)]
    -> [(PackageName, ModuleName, FixityInfo)]
    forall a. Semigroup a => a -> a -> a
    (<>)
        ([(OpName, [(PackageName, ModuleName, FixityInfo)])]
     -> Map OpName [(PackageName, ModuleName, FixityInfo)])
    -> (Map PackageName (Map ModuleName (Map OpName FixityInfo))
        -> [(OpName, [(PackageName, ModuleName, FixityInfo)])])
    -> Map PackageName (Map ModuleName (Map OpName FixityInfo))
    -> Map OpName [(PackageName, ModuleName, FixityInfo)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((PackageName, (ModuleName, (OpName, FixityInfo)))
     -> (OpName, [(PackageName, ModuleName, FixityInfo)]))
    -> [(PackageName, (ModuleName, (OpName, FixityInfo)))]
    -> [(OpName, [(PackageName, ModuleName, FixityInfo)])]
    forall a b. (a -> b) -> [a] -> [b]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (PackageName, (ModuleName, (OpName, FixityInfo)))
    -> (OpName, [(PackageName, ModuleName, FixityInfo)])
    forall {a} {b} {a} {c}. (a, (b, (a, c))) -> (a, [(a, b, c)])
    rearrange
        ([(PackageName, (ModuleName, (OpName, FixityInfo)))]
     -> [(OpName, [(PackageName, ModuleName, FixityInfo)])])
    -> (Map PackageName (Map ModuleName (Map OpName FixityInfo))
        -> [(PackageName, (ModuleName, (OpName, FixityInfo)))])
    -> Map PackageName (Map ModuleName (Map OpName FixityInfo))
    -> [(OpName, [(PackageName, ModuleName, FixityInfo)])]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [(PackageName, [(ModuleName, (OpName, FixityInfo))])]
    -> [(PackageName, (ModuleName, (OpName, FixityInfo)))]
    forall {m :: * -> *} {a} {b}. Monad m => m (a, m b) -> m (a, b)
    flatten
        ([(PackageName, [(ModuleName, (OpName, FixityInfo))])]
     -> [(PackageName, (ModuleName, (OpName, FixityInfo)))])
    -> (Map PackageName (Map ModuleName (Map OpName FixityInfo))
        -> [(PackageName, [(ModuleName, (OpName, FixityInfo))])])
    -> Map PackageName (Map ModuleName (Map OpName FixityInfo))
    -> [(PackageName, (ModuleName, (OpName, FixityInfo)))]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Map PackageName [(ModuleName, (OpName, FixityInfo))]
    -> [(PackageName, [(ModuleName, (OpName, FixityInfo))])]
    forall k a. Map k a -> [(k, a)]
    Map.toList
        (Map PackageName [(ModuleName, (OpName, FixityInfo))]
     -> [(PackageName, [(ModuleName, (OpName, FixityInfo))])])
    -> (Map PackageName (Map ModuleName (Map OpName FixityInfo))
        -> Map PackageName [(ModuleName, (OpName, FixityInfo))])
    -> Map PackageName (Map ModuleName (Map OpName FixityInfo))
    -> [(PackageName, [(ModuleName, (OpName, FixityInfo))])]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Map ModuleName (Map OpName FixityInfo)
     -> [(ModuleName, (OpName, FixityInfo))])
    -> Map PackageName (Map ModuleName (Map OpName FixityInfo))
    -> Map PackageName [(ModuleName, (OpName, FixityInfo))]
    forall a b k. (a -> b) -> Map k a -> Map k b
    Map.map ([(ModuleName, [(OpName, FixityInfo)])]
    -> [(ModuleName, (OpName, FixityInfo))]
    forall {m :: * -> *} {a} {b}. Monad m => m (a, m b) -> m (a, b)
    flatten ([(ModuleName, [(OpName, FixityInfo)])]
     -> [(ModuleName, (OpName, FixityInfo))])
    -> (Map ModuleName (Map OpName FixityInfo)
        -> [(ModuleName, [(OpName, FixityInfo)])])
    -> Map ModuleName (Map OpName FixityInfo)
    -> [(ModuleName, (OpName, FixityInfo))]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Map ModuleName [(OpName, FixityInfo)]
    -> [(ModuleName, [(OpName, FixityInfo)])]
    forall k a. Map k a -> [(k, a)]
    Map.toList (Map ModuleName [(OpName, FixityInfo)]
     -> [(ModuleName, [(OpName, FixityInfo)])])
    -> (Map ModuleName (Map OpName FixityInfo)
        -> Map ModuleName [(OpName, FixityInfo)])
    -> Map ModuleName (Map OpName FixityInfo)
    -> [(ModuleName, [(OpName, FixityInfo)])]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Map OpName FixityInfo -> [(OpName, FixityInfo)])
    -> Map ModuleName (Map OpName FixityInfo)
    -> Map ModuleName [(OpName, FixityInfo)]
    forall a b k. (a -> b) -> Map k a -> Map k b
    Map.map Map OpName FixityInfo -> [(OpName, FixityInfo)]
    forall k a. Map k a -> [(k, a)]
    Map.toList)
        (Map PackageName (Map ModuleName (Map OpName FixityInfo))
     -> PackageFixityMap)
    -> Map PackageName (Map ModuleName (Map OpName FixityInfo))
    -> PackageFixityMap
    forall a b. (a -> b) -> a -> b
    $ Map PackageName (Map ModuleName (Map OpName FixityInfo))
    -> Set PackageName
    -> Map PackageName (Map ModuleName (Map OpName FixityInfo))
    forall k a. Ord k => Map k a -> Set k -> Map k a
    Map.restrictKeys Map PackageName (Map ModuleName (Map OpName FixityInfo))
    m Set PackageName
    dependencies
      where
        rearrange :: (a, (b, (a, c))) -> (a, [(a, b, c)])
    rearrange (a
    packageName, (b
    moduleName, (a
    opName, c
    fixityInfo))) =
          (a
    opName, [(a
    packageName, b
    moduleName, c
    fixityInfo)])
        flatten :: m (a, m b) -> m (a, b)
    flatten m (a, m b)
    xs = do
          (k, vs) <- m (a, m b)
    xs
          v <- vs
          return (k, v)

    -- | Compute the fixity map that is specific to the module we are formatting.
    moduleFixityMap ::
      -- | Fixity information selected from dependencies of this package
      PackageFixityMap ->
      -- | A simplified representation of the import list in this module
      [FixityImport] ->
      -- | Fixity map specific to this module
      ModuleFixityMap
    moduleFixityMap :: PackageFixityMap -> [FixityImport] -> ModuleFixityMap
    moduleFixityMap (PackageFixityMap Map OpName (NonEmpty (PackageName, ModuleName, FixityInfo))
    m) [FixityImport]
    imports =
      Map OpName FixityProvenance -> ModuleFixityMap
    ModuleFixityMap (Map OpName FixityProvenance -> ModuleFixityMap)
    -> Map OpName FixityProvenance -> ModuleFixityMap
    forall a b. (a -> b) -> a -> b
    $
        OpName
    -> FixityProvenance
    -> Map OpName FixityProvenance
    -> Map OpName FixityProvenance
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    Map.insert
          OpName
    ":"
          (FixityInfo -> FixityProvenance
    Given FixityInfo
    colonFixityInfo)
          ((NonEmpty (FixityQualification, FixityInfo) -> FixityProvenance)
    -> Map OpName (NonEmpty (FixityQualification, FixityInfo))
    -> Map OpName FixityProvenance
    forall a b k. (a -> b) -> Map k a -> Map k b
    Map.map NonEmpty (FixityQualification, FixityInfo) -> FixityProvenance
    FromModuleImports ((OpName
     -> NonEmpty (PackageName, ModuleName, FixityInfo)
     -> Maybe (NonEmpty (FixityQualification, FixityInfo)))
    -> Map OpName (NonEmpty (PackageName, ModuleName, FixityInfo))
    -> Map OpName (NonEmpty (FixityQualification, FixityInfo))
    forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
    Map.mapMaybeWithKey OpName
    -> NonEmpty (PackageName, ModuleName, FixityInfo)
    -> Maybe (NonEmpty (FixityQualification, FixityInfo))
    select Map OpName (NonEmpty (PackageName, ModuleName, FixityInfo))
    m))
      where
        select ::
          OpName ->
          NonEmpty (PackageName, ModuleName, FixityInfo) ->
          Maybe (NonEmpty (FixityQualification, FixityInfo))
        select :: OpName
    -> NonEmpty (PackageName, ModuleName, FixityInfo)
    -> Maybe (NonEmpty (FixityQualification, FixityInfo))
    select OpName
    opName =
          let f :: (PackageName, ModuleName, t) -> [(FixityQualification, t)]
    f (PackageName
    packageName, ModuleName
    moduleName, t
    fixityInfo) =
                (,t
    fixityInfo)
                  (FixityQualification -> (FixityQualification, t))
    -> [FixityQualification] -> [(FixityQualification, t)]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> PackageName -> ModuleName -> OpName -> [FixityQualification]
    resolveThroughImports PackageName
    packageName ModuleName
    moduleName OpName
    opName
           in [(FixityQualification, FixityInfo)]
    -> Maybe (NonEmpty (FixityQualification, FixityInfo))
    forall a. [a] -> Maybe (NonEmpty a)
    NE.nonEmpty ([(FixityQualification, FixityInfo)]
     -> Maybe (NonEmpty (FixityQualification, FixityInfo)))
    -> (NonEmpty (PackageName, ModuleName, FixityInfo)
        -> [(FixityQualification, FixityInfo)])
    -> NonEmpty (PackageName, ModuleName, FixityInfo)
    -> Maybe (NonEmpty (FixityQualification, FixityInfo))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((PackageName, ModuleName, FixityInfo)
     -> [(FixityQualification, FixityInfo)])
    -> NonEmpty (PackageName, ModuleName, FixityInfo)
    -> [(FixityQualification, FixityInfo)]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (PackageName, ModuleName, FixityInfo)
    -> [(FixityQualification, FixityInfo)]
    forall {t}.
    (PackageName, ModuleName, t) -> [(FixityQualification, t)]
    f
        resolveThroughImports ::
          PackageName ->
          ModuleName ->
          OpName ->
          [FixityQualification]
        resolveThroughImports :: PackageName -> ModuleName -> OpName -> [FixityQualification]
    resolveThroughImports PackageName
    packageName ModuleName
    moduleName OpName
    opName =
          let doesImportMatch :: FixityImport -> Bool
    doesImportMatch FixityImport {Maybe (ImportListInterpretation, [OpName])
    Maybe PackageName
    ModuleName
    FixityQualification
    fimportPackage :: Maybe PackageName
    fimportModule :: ModuleName
    fimportQualified :: FixityQualification
    fimportList :: Maybe (ImportListInterpretation, [OpName])
    fimportList :: FixityImport -> Maybe (ImportListInterpretation, [OpName])
    fimportQualified :: FixityImport -> FixityQualification
    fimportModule :: FixityImport -> ModuleName
    fimportPackage :: FixityImport -> Maybe PackageName
    ..} =
                let packageMatches :: Bool
    packageMatches =
                      case Maybe PackageName
    fimportPackage of
                        Maybe PackageName
    Nothing -> Bool
    True
                        Just PackageName
    p -> PackageName
    p PackageName -> PackageName -> Bool
    forall a. Eq a => a -> a -> Bool
    == PackageName
    packageName
                    moduleMatches :: Bool
    moduleMatches =
                      ModuleName
    fimportModule ModuleName -> ModuleName -> Bool
    forall a. Eq a => a -> a -> Bool
    == ModuleName
    moduleName
                    opMatches :: Bool
    opMatches = case Maybe (ImportListInterpretation, [OpName])
    fimportList of
                      Maybe (ImportListInterpretation, [OpName])
    Nothing -> Bool
    True
                      Just (ImportListInterpretation
    Exactly, [OpName]
    xs) -> OpName
    opName OpName -> [OpName] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `elem` [OpName]
    xs
                      Just (ImportListInterpretation
    EverythingBut, [OpName]
    xs) -> OpName
    opName OpName -> [OpName] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `notElem` [OpName]
    xs
                 in Bool
    packageMatches Bool -> Bool -> Bool
    && Bool
    moduleMatches Bool -> Bool -> Bool
    && Bool
    opMatches
           in FixityImport -> FixityQualification
    fimportQualified (FixityImport -> FixityQualification)
    -> [FixityImport] -> [FixityQualification]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (FixityImport -> Bool) -> [FixityImport] -> [FixityImport]
    forall a. (a -> Bool) -> [a] -> [a]
    filter FixityImport -> Bool
    doesImportMatch [FixityImport]
    imports

    -- | Apply fixity overrides.
    applyFixityOverrides ::
      -- | User overrides
      FixityOverrides ->
      -- | Module fixity map
      ModuleFixityMap ->
      -- | Module fixity map with overrides applied
      ModuleFixityMap
    applyFixityOverrides :: FixityOverrides -> ModuleFixityMap -> ModuleFixityMap
    applyFixityOverrides (FixityOverrides Map OpName FixityInfo
    o) (ModuleFixityMap Map OpName FixityProvenance
    m) =
      Map OpName FixityProvenance -> ModuleFixityMap
    ModuleFixityMap (Map OpName FixityProvenance
    -> Map OpName FixityProvenance -> Map OpName FixityProvenance
    forall k a. Ord k => Map k a -> Map k a -> Map k a
    Map.union ((FixityInfo -> FixityProvenance)
    -> Map OpName FixityInfo -> Map OpName FixityProvenance
    forall a b k. (a -> b) -> Map k a -> Map k b
    Map.map FixityInfo -> FixityProvenance
    Given Map OpName FixityInfo
    o) Map OpName FixityProvenance
    m)

    memoSet :: (Set PackageName -> v) -> Set PackageName -> v
    memoSet :: forall v. (Set PackageName -> v) -> Set PackageName -> v
    memoSet Set PackageName -> v
    f =
      ([String] -> v) -> [String] -> v
    forall t a. HasTrie t => (t -> a) -> t -> a
    memo (Set PackageName -> v
    f (Set PackageName -> v)
    -> ([String] -> Set PackageName) -> [String] -> v
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [PackageName] -> Set PackageName
    forall a. Eq a => [a] -> Set a
    Set.fromAscList ([PackageName] -> Set PackageName)
    -> ([String] -> [PackageName]) -> [String] -> Set PackageName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (String -> PackageName) -> [String] -> [PackageName]
    forall a b. (a -> b) -> [a] -> [b]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap String -> PackageName
    mkPackageName)
        ([String] -> v)
    -> (Set PackageName -> [String]) -> Set PackageName -> v
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (PackageName -> String) -> [PackageName] -> [String]
    forall a b. (a -> b) -> [a] -> [b]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap PackageName -> String
    unPackageName
        ([PackageName] -> [String])
    -> (Set PackageName -> [PackageName])
    -> Set PackageName
    -> [String]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Set PackageName -> [PackageName]
    forall a. Set a -> [a]
    Set.toAscList