{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}

module Ormolu.Config.TH
  ( allNothing,
    unpackFieldsWithSuffix,

    -- * BijectiveMap
    BijectiveMap,
    mkBijectiveMap,
    parseTextWith,
    showTextWith,
    showAllValues,
  )
where

import Control.Monad (forM, when, (>=>))
import Data.Containers.ListUtils (nubOrd)
import Data.List (intercalate, nub)
import Language.Haskell.TH
import Language.Haskell.TH.Syntax (lift)
import Text.Printf (printf)

allNothing :: Name -> Q Exp
allNothing :: Name -> Q Exp
allNothing Name
name = do
  Type
ty <- Name -> Q Type
reifyType Name
name
  forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
name) forall a b. (a -> b) -> a -> b
$
    forall a. Int -> a -> [a]
replicate (Type -> Int
getArity Type
ty) [|Nothing|]

unpackFieldsWithSuffix :: Name -> String -> Q Pat
unpackFieldsWithSuffix :: Name -> String -> Q Pat
unpackFieldsWithSuffix Name
name String
suffix = do
  Name
typeForCon <-
    Name -> Q Info
reify Name
name forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      DataConI Name
_ Type
_ Name
typeForCon -> forall (m :: * -> *) a. Monad m => a -> m a
return Name
typeForCon
      Info
info -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"allNothing requires the Name of a data constructor, got: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Info
info

  [Con]
allConsInType <-
    Name -> Q (Either Info [Con])
getAllConstructors Name
typeForCon
      forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r. PrintfType r => String -> r
printf String
"Unexpected parent of data constructor: %s" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show) forall (m :: * -> *) a. Monad m => a -> m a
return

  [Name]
fields <-
    case forall a. (a -> Bool) -> [a] -> [a]
filter (forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem Name
name forall b c a. (b -> c) -> (a -> b) -> a -> c
. Con -> [Name]
getConstructorNames) [Con]
allConsInType of
      [Con
con] | Just [Name]
fields <- Con -> Maybe [Name]
conFieldNames Con
con -> forall (m :: * -> *) a. Monad m => a -> m a
return [Name]
fields
      [Con]
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"Could not find unique record constructor in: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show [Con]
allConsInType

  forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
name forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (m :: * -> *). Quote m => Name -> m Pat
varP forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Name
mkName forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Semigroup a => a -> a -> a
<> String
suffix) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase) [Name]
fields
  where
    conFieldNames :: Con -> Maybe [Name]
conFieldNames = \case
      NormalC {} -> forall a. Maybe a
Nothing
      RecC Name
_ [VarBangType]
tys -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall {a} {b} {c}. (a, b, c) -> a
fst3 [VarBangType]
tys
      InfixC {} -> forall a. Maybe a
Nothing
      ForallC {} -> forall a. Maybe a
Nothing
      GadtC {} -> forall a. Maybe a
Nothing
      RecGadtC [Name]
_ [VarBangType]
tys Type
_ -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall {a} {b} {c}. (a, b, c) -> a
fst3 [VarBangType]
tys
    fst3 :: (a, b, c) -> a
fst3 (a
x, b
_, c
_) = a
x

data BijectiveMap a = BijectiveMap
  { forall a. BijectiveMap a -> String -> Either String a
parseTextWith :: String -> Either String a,
    forall a. BijectiveMap a -> a -> String
showTextWith :: a -> String,
    forall a. BijectiveMap a -> [String]
getAllOptions :: [String]
  }

showAllValues :: BijectiveMap a -> String
showAllValues :: forall a. BijectiveMap a -> String
showAllValues = [String] -> String
uncommas forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. BijectiveMap a -> [String]
getAllOptions

-- | Generate a `BijectiveMap a` value with the given map.
--
-- Checks the following:
--   * all Names in given list refer to a constructor of type `a`
--   * all Names in given list refer to a 0-arity constructor
--   * all constructors in type `a` are accounted for.
--   * each constructor in type `a` must be provided only once.
mkBijectiveMap :: [(Name, String)] -> Q Exp
mkBijectiveMap :: [(Name, String)] -> Q Exp
mkBijectiveMap [(Name, String)]
mapping = do
  let ([Name]
conNames, [String]
allOptions) = forall a b. [(a, b)] -> ([a], [b])
unzip [(Name, String)]
mapping

  -- check all names refer to constructors
  ([Type]
conTypes, [Name]
conParents) <-
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. [(a, b)] -> ([a], [b])
unzip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Name]
conNames forall a b. (a -> b) -> a -> b
$ \Name
name ->
      Name -> Q Info
reify Name
name forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        DataConI Name
_ Type
ty Name
parent -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
ty, Name
parent)
        Info
info ->
          forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$
            forall r. PrintfType r => String -> r
printf
              String
"mkBijectiveMap requires all Names refer to data constructors, got %s: %s"
              (forall a. Show a => a -> String
show Name
name)
              (forall a. Show a => a -> String
show Info
info)

  -- check that all constructors are in same type
  Name
parent <-
    case forall a. Eq a => [a] -> [a]
nub [Name]
conParents of
      [Name
parent] -> forall (m :: * -> *) a. Monad m => a -> m a
return Name
parent
      [Name]
parents -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"mkBijectiveMap requires all Names refer to data constructors in the same type, got: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show [Name]
parents

  -- check that all constructors are 0-arity
  case forall a. (a -> Bool) -> [a] -> [a]
filter ((forall a. Eq a => a -> a -> Bool
/= Int
0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Int
getArity) [Type]
conTypes of
    [] -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
    [Type]
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"mkBijectiveMap requires all constructors have 0-arity"

  -- check that all constructors are represented
  [Con]
allConsInType <-
    Name -> Q (Either Info [Con])
getAllConstructors Name
parent
      forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r. PrintfType r => String -> r
printf String
"Unexpected parent of data constructors: %s" forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show) forall (m :: * -> *) a. Monad m => a -> m a
return
  case forall a. (a -> Bool) -> [a] -> [a]
filter (forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Name]
conNames) (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Con -> [Name]
getConstructorNames [Con]
allConsInType) of
    [] -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
    [Name]
missing -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"Missing constructors: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show [Name]
missing

  -- check for duplicate constructors
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Ord a => [a] -> [a]
nubOrd [Name]
conNames forall a. Eq a => a -> a -> Bool
/= [Name]
conNames) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"mkBijectiveMap requires each constructor to be provided only once"

  Name
unknown <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"unknown"
  let parser :: Q Exp
parser =
        forall (m :: * -> *). Quote m => [m Match] -> m Exp
lamCaseE forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$
          [ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> [a] -> [b]
map [(Name, String)]
mapping forall a b. (a -> b) -> a -> b
$ \(Name
name, String
option) ->
              forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match
                (forall (m :: * -> *). Quote m => Lit -> m Pat
litP forall a b. (a -> b) -> a -> b
$ String -> Lit
stringL String
option)
                (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|Right $(conE name)|])
                [],
            [ forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match
                (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
unknown)
                ( forall (m :: * -> *). Quote m => m Exp -> m Body
normalB
                    [|
                      Left . unlines $
                        [ "unknown value: " <> show $(varE unknown),
                          "Valid values are: " <> $(lift $ uncommas $ map show allOptions)
                        ]
                      |]
                )
                []
            ]
          ]
      shower :: Q Exp
shower =
        forall (m :: * -> *). Quote m => [m Match] -> m Exp
lamCaseE forall a b. (a -> b) -> a -> b
$
          forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> [a] -> [b]
map [(Name, String)]
mapping forall a b. (a -> b) -> a -> b
$ \(Name
name, String
option) ->
            forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
name []) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB forall a b. (a -> b) -> a -> b
$ forall t (m :: * -> *). (Lift t, Quote m) => t -> m Exp
lift String
option) []

  [|
    BijectiveMap
      { parseTextWith = $parser,
        showTextWith = $shower,
        getAllOptions = $(lift allOptions)
      }
    |]

----------------------------------------------------------------------------
-- Helpers

{- FOURMOLU_DISABLE -}
{- https://github.com/fourmolu/fourmolu#limitations -}
getArity :: Type -> Int
getArity :: Type -> Int
getArity = \case
  ForallT [TyVarBndr Specificity]
_ [Type]
_ Type
ty -> Type -> Int
getArity Type
ty
  AppT (AppT Type
ArrowT Type
_) Type
ty -> Int
1 forall a. Num a => a -> a -> a
+ Type -> Int
getArity Type
ty
#if MIN_VERSION_template_haskell(2,17,0)
  AppT (AppT (AppT Type
MulArrowT Type
_) Type
_) Type
ty -> Int
1 forall a. Num a => a -> a -> a
+ Type -> Int
getArity Type
ty
#endif
  Type
_ -> Int
0
{- FOURMOLU_ENABLE -}

getAllConstructors :: Name -> Q (Either Info [Con])
getAllConstructors :: Name -> Q (Either Info [Con])
getAllConstructors =
  Name -> Q Info
reify forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> \case
    TyConI (DataD [Type]
_ Name
_ [TyVarBndr ()]
_ Maybe Type
_ [Con]
cons [DerivClause]
_) -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right [Con]
cons
    Info
info -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left Info
info

-- Could return multiple names for GADTs like 'A, B :: Foo'
getConstructorNames :: Con -> [Name]
getConstructorNames :: Con -> [Name]
getConstructorNames = \case
  NormalC Name
n [BangType]
_ -> [Name
n]
  RecC Name
n [VarBangType]
_ -> [Name
n]
  InfixC BangType
_ Name
n BangType
_ -> [Name
n]
  ForallC [TyVarBndr Specificity]
_ [Type]
_ Con
c -> Con -> [Name]
getConstructorNames Con
c
  GadtC [Name]
ns [BangType]
_ Type
_ -> [Name]
ns
  RecGadtC [Name]
ns [VarBangType]
_ Type
_ -> [Name]
ns

uncommas :: [String] -> String
uncommas :: [String] -> String
uncommas [] = String
""
uncommas [String
s] = String
s
uncommas [String
s0, String
s1] = String
s0 forall a. Semigroup a => a -> a -> a
<> String
" or " forall a. Semigroup a => a -> a -> a
<> String
s1
uncommas [String]
ss =
  let pre :: [String]
pre = forall a. [a] -> [a]
init [String]
ss
      end :: String
end = forall a. [a] -> a
last [String]
ss
   in forall a. [a] -> [[a]] -> [a]
intercalate String
", " [String]
pre forall a. Semigroup a => a -> a -> a
<> String
"or " forall a. Semigroup a => a -> a -> a
<> String
end