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

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

  , GenericTypeName
  ) where

import Control.Exception (assert)
import Data.Vector qualified as V
import GHC.Generics qualified as G
import GHC.TypeLits (Symbol)
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 :: forall a. (Natural -> a -> a -> a) -> NonEmpty a -> a
mkGenericTree Natural -> a -> a -> a
mkNode = (a -> a) -> (Natural -> a -> a -> a) -> Vector a -> a
forall a b.
HasCallStack =>
(a -> b) -> (Natural -> b -> b -> b) -> Vector a -> b
mkGenericTreeVec a -> a
forall a. a -> a
id Natural -> a -> a -> a
mkNode (Vector a -> a) -> (NonEmpty a -> Vector a) -> NonEmpty a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Vector a
forall a. [a] -> Vector a
V.fromList ([a] -> Vector a) -> (NonEmpty a -> [a]) -> NonEmpty a -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty a -> [a]
forall t. Container t => t -> [Element t]
toList

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

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

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