Shared.hs

{-# LANGUAGE UndecidableInstances #-}
-- | Additional classes that help in comparing values in tests.
module Shared
  ( lowercase, HasShape (shapeL), Linearizable (linearize)
  ) where

import Prelude

import Data.Array.Internal.RankedS qualified as RS
import Data.Char qualified
import Data.Foldable qualified as Foldable
import Data.Int (Int16, Int32, Int64, Int8)
import Data.Vector.Storable qualified as VS
import Foreign.C (CInt)

import Data.Array.Nested qualified as Nested
import Data.Array.Nested.Mixed qualified as Nested.Mixed
import Data.Array.Nested.Ranked.Shape
import Data.Array.Nested.Shaped (stoPrimitive)
import Data.Array.Nested.Shaped.Shape
import Data.Array.XArray qualified as X

import HordeAd.Core.CarriersConcrete
import HordeAd.Core.Types

lowercase :: String -> String
lowercase = map Data.Char.toLower

-- | Things that have shape.
class HasShape a where
  shapeL :: a -> [Int]

instance Nested.Elt a => HasShape (Nested.Ranked n a) where
  shapeL = shrToList . Nested.rshape

instance Nested.Elt a => HasShape (Nested.Shaped sh a) where
  shapeL = shsToList . Nested.sshape

instance HasShape (RepConcrete y) => HasShape (Concrete y) where
  shapeL = shapeL . unConcrete

instance HasShape Double where
  shapeL _ = []

instance HasShape Float where
  shapeL _ = []

instance HasShape Int64 where
  shapeL _ = []

instance HasShape Int32 where
  shapeL _ = []

instance HasShape Int16 where
  shapeL _ = []

instance HasShape Int8 where
  shapeL _ = []

instance HasShape Int where
  shapeL _ = []

instance HasShape CInt where
  shapeL _ = []

instance HasShape Z1 where
  shapeL _ = [0]

instance {-# OVERLAPPABLE #-} (Foldable t) => HasShape (t a) where
  shapeL = (: []) . length

-- | Things that can be linearized, i.e. converted to a list.
class Linearizable a b | a -> b where
  linearize :: a -> [b]

instance (VS.Storable a, Nested.PrimElt a)
         => Linearizable (Nested.Ranked n a) a where
  linearize = RS.toList . Nested.rtoOrthotope

instance (VS.Storable a, Nested.PrimElt a)
         => Linearizable (Nested.Shaped sh a) a where
  linearize a =
    let Nested.Shaped (Nested.Mixed.M_Primitive _ (X.XArray arr)) =
          stoPrimitive a
    in RS.toList arr

instance Linearizable (RepConcrete y) a
         => Linearizable (Concrete y) a where
  linearize = linearize . unConcrete

instance Linearizable Double Double where
  linearize x = [x]

instance Linearizable Float Float where
  linearize x = [x]

instance Linearizable Int64 Int64 where
  linearize x = [x]

instance Linearizable Int32 Int32 where
  linearize x = [x]

instance Linearizable Int16 Int16 where
  linearize x = [x]

instance Linearizable Int8 Int8 where
  linearize x = [x]

instance Linearizable Int Int where
  linearize x = [x]

instance Linearizable CInt CInt where
  linearize x = [x]

instance Linearizable Z1 Z1 where
  linearize _ = []

instance {-# OVERLAPPABLE #-} (Foldable t) => Linearizable (t a) a where
  linearize = Foldable.toList