-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Generic-related utils.
module Morley.Util.Generic
  ( mkGenericTree
  , mkGenericTreeVec

  , GenericTypeName

  , GRep
  , NiceGeneric
  ) where

import Control.Exception (assert)
import Data.Vector qualified as V
import GHC.Generics qualified as G
import GHC.TypeLits (Symbol)
import Type.Errors (ErrorMessage(..), IfStuck, Pure, TypeError)
import Unsafe qualified (fromIntegral)

-- | Rebuild a list into a binary tree of exactly the same form which
-- "Data.Generics" uses to represent datatypes.
--
-- Along with the original list you have to provide constructor for intermediate
-- nodes - it accepts zero-based index of the leftmost element of the right tree
-- and merged trees themselves.
mkGenericTree :: (Natural -> a -> a -> a) -> NonEmpty a -> a
mkGenericTree mkNode = mkGenericTreeVec id mkNode . V.fromList . toList

mkGenericTreeVec
  :: HasCallStack
  => (a -> b) -> (Natural -> b -> b -> b) -> V.Vector a -> b
mkGenericTreeVec mkLeaf mkNode vector
  | V.null vector = error "Empty vector"
  | otherwise = mkTreeDo 0 vector
  where
    mkTreeDo idxBase es
      | V.length es == 1 = mkLeaf $ V.head es
      | otherwise = assert (V.length es > 1) $
          let mid = V.length es `div` 2
              mid' = idxBase + mid
              (h, t) = V.splitAt mid es
          in mkNode (Unsafe.fromIntegral @Int @Natural mid') (mkTreeDo idxBase h) (mkTreeDo mid' t)

-- | Extract datatype name via its Generic representation.
--
-- For polymorphic types this throws away all type arguments.
type GenericTypeName a = GTypeName (GRep a)

type family GTypeName (x :: Type -> Type) :: Symbol where
  GTypeName (G.D1 ('G.MetaData tyName _ _ _) _) = tyName

-- | Avoid too eagerly reducing 'GRep'.
data ThisTypeShallNotBeExported

{- | Like 'G.Rep', but has better error messages when stuck.

Trying to use something requiring generics without the 'Generic' instance should
trigger a custom error message:

>>> data Foo = Foo
>>> (from . to :: GRep a ~ GRep () => a -> ()) Foo
...
... GHC.Generics.Rep Foo
... is stuck. Likely
... Generic Foo
... instance is missing or out of scope.
...
>>> data Foo = Foo deriving Generic
>>> (from . to :: GRep a ~ GRep () => a -> ()) Foo
...
... Couldn't match type ‘"Foo"’ with ‘"()"’
...
-}
type family GRep t where
  GRep ThisTypeShallNotBeExported = TypeError ('Text "impossible")
  GRep t = IfStuck (G.Rep t) (TypeError (GRepErrorMsg t)) (Pure (G.Rep t))

type GRepErrorMsg t = 'ShowType (G.Rep t)
  ':$$: 'Text "is stuck. Likely"
  ':$$: 'ShowType (Generic t)
  ':$$: 'Text "instance is missing or out of scope."

-- | Similar to 'Generic', but also asserts that 'GRep' is the same as 'G.Rep'.
-- Useful in place of 'Generic' when 'GRep' is used.
type NiceGeneric t = (Generic t, GRep t ~ G.Rep t)