Data.Ord.Tuple

Tuple.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DerivingVia #-}

{-# OPTIONS_GHC -ddump-simpl #-}
{-# OPTIONS_GHC -dsuppress-all #-}
{-# OPTIONS_GHC -dno-suppress-type-signatures #-}
{-# OPTIONS_GHC -ddump-to-file #-}

{-# OPTIONS_GHC -fspecialise-aggressively #-}
{-# OPTIONS_GHC -fdicts-cheap #-}
{-# OPTIONS_GHC -fdicts-strict #-}
{-# OPTIONS_GHC -fmax-worker-args=100 #-}

module Data.Ord.Tuple
  ( OrdTuple (..)
  , OrderBy (..)
  , HasComponent (..)
  , sortBySpec
  ) where

import Data.Kind
import Data.Proxy
import Data.Tuple
import GHC.Exts (Proxy#, proxy#)
import GHC.List (List)
import GHC.TypeNats

import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as MG
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as MU
import GHC.Generics (Generic)

import Data.Vector.Algorithms.Quicksort qualified as Quick
import Data.Coerce
import Control.Monad.ST
import Data.Vector.Algorithms.Quicksort.Parameterised

import qualified Data.Strict.Tuple as Strict

data OrderBy = Up Nat | Down Nat

newtype OrdTuple (ixs :: List OrderBy) (a :: Type) = OrdTuple a

newtype instance U.MVector s (OrdTuple ixs a) = MV_OrdTuple (U.MVector s a)
newtype instance U.Vector    (OrdTuple ixs a) = V_OrdTuple  (U.Vector   a)
deriving instance MG.MVector U.MVector a => MG.MVector U.MVector (OrdTuple ixs a)
deriving instance G.Vector U.Vector a => G.Vector U.Vector (OrdTuple ixs a)
instance U.Unbox a => U.Unbox (OrdTuple ixs a)



newtype instance U.MVector s (Strict.Pair a b) = MV_StrictPair (U.MVector s (a, b))
newtype instance U.Vector    (Strict.Pair a b) = V_StrictPair  (U.Vector    (a, b))
instance U.IsoUnbox (Strict.Pair a b) (a, b)
deriving via (Strict.Pair a b `U.As` (a, b)) instance (U.Unbox a, U.Unbox b)
  => MG.MVector U.MVector (Strict.Pair a b)
deriving via (Strict.Pair a b `U.As` (a, b)) instance (U.Unbox a, U.Unbox b)
  => G.Vector  U.Vector  (Strict.Pair a b)
instance (U.Unbox a, U.Unbox b) => U.Unbox (Strict.Pair a b)


instance Eq (OrdTuple '[] a) where
  (==) = const $ const True

instance Ord (OrdTuple '[] a) where
  compare = const $ const EQ
  (<=) = const $ const True
  (>=) = const $ const True
  (<) = const $ const False
  (>) = const $ const False

instance {-# OVERLAPPING #-}
  (HasComponent ix a b, Eq b)
  => Eq (OrdTuple (dir ix ': '[]) a)
  where
  OrdTuple a == OrdTuple a' =
    getComponent (proxy# @ix) a == getComponent (proxy# @ix) a'

instance {-# OVERLAPPING #-}
  (HasComponent ix a b, Ord b)
  => Ord (OrdTuple ('Up ix ': '[]) a)
  where
  compare (OrdTuple a) (OrdTuple a') =
    compare (getComponent (proxy# @ix) a) (getComponent (proxy# @ix) a')
  OrdTuple a <= OrdTuple a' =
    getComponent (proxy# @ix) a <= getComponent (proxy# @ix) a'
  OrdTuple a >= OrdTuple a' =
    getComponent (proxy# @ix) a <= getComponent (proxy# @ix) a'
  OrdTuple a < OrdTuple a' =
    getComponent (proxy# @ix) a <= getComponent (proxy# @ix) a'
  OrdTuple a > OrdTuple a' =
    getComponent (proxy# @ix) a <= getComponent (proxy# @ix) a'

instance {-# OVERLAPPING #-}
  (HasComponent ix a b, Ord b)
  => Ord (OrdTuple ('Down ix ': '[]) a)
  where
  compare (OrdTuple a) (OrdTuple a') =
    compare (getComponent (proxy# @ix) a') (getComponent (proxy# @ix) a)
  OrdTuple a <= OrdTuple a' =
    getComponent (proxy# @ix) a' <= getComponent (proxy# @ix) a
  OrdTuple a >= OrdTuple a' =
    getComponent (proxy# @ix) a' <= getComponent (proxy# @ix) a
  OrdTuple a < OrdTuple a' =
    getComponent (proxy# @ix) a' <= getComponent (proxy# @ix) a
  OrdTuple a > OrdTuple a' =
    getComponent (proxy# @ix) a' <= getComponent (proxy# @ix) a

instance
  (HasComponent ix a b, Eq b, Eq (OrdTuple ixs a))
  => Eq (OrdTuple (dir ix ': ixs) a)
  where
  OrdTuple a == OrdTuple a' =
    getComponent (proxy# @ix) a == getComponent (proxy# @ix) a'
      && (OrdTuple @ixs a) == OrdTuple a'

instance
  (HasComponent ix a b, Ord b, Ord (OrdTuple ixs a))
  => Ord (OrdTuple ('Up ix ': ixs) a)
  where
  compare (OrdTuple a) (OrdTuple a') =
    compare (getComponent (proxy# @ix) a) (getComponent (proxy# @ix) a')
      <> compare (OrdTuple @ixs a) (OrdTuple a')

instance
  (HasComponent ix a b, Ord b, Ord (OrdTuple ixs a))
  => Ord (OrdTuple ('Down ix ': ixs) a)
  where
  compare (OrdTuple a) (OrdTuple a') =
    compare (getComponent (proxy# @ix) a') (getComponent (proxy# @ix) a)
      <> compare (OrdTuple @ixs a) (OrdTuple a')

class HasComponent (ix :: Nat) a b | ix a -> b where
  getComponent :: Proxy# ix -> a -> b

instance HasComponent 0 (Solo a) a where
  getComponent _ (MkSolo a) = a

instance HasComponent 0 (a, b) a where
  getComponent _ (a, _) = a

instance HasComponent 0 (Strict.Pair a b) a where
  getComponent _ (a Strict.:!: _) = a

instance HasComponent 1 (Strict.Pair a b) b where
  getComponent _ (_ Strict.:!: b) = b


instance HasComponent 0 (a, b, c) a where
  getComponent _ (a, _, _) = a

instance HasComponent 1 (a, b) b where
  getComponent _ (_, b) = b

instance HasComponent 1 (a, b, c) b where
  getComponent _ (_, b, _) = b

instance HasComponent 2 (a, b, c) c where
  getComponent _ (_, _, c) = c













sortBy
  :: forall v a b.
     (Ord b, G.Vector v b, Coercible (v a) (v b))
  => Proxy b
  -> v a
  -> v a
sortBy _ vs = coerce (Quick.sort (coerce vs :: v b))

sortByOrdTuple
  :: forall v a (ixs :: [OrderBy]).
     (Ord (OrdTuple ixs a), G.Vector v (OrdTuple ixs a), Coercible (v a) (v (OrdTuple ixs a)))
  => Proxy ixs
  -> v a
  -> v a
sortByOrdTuple _ = sortBy (Proxy :: Proxy (OrdTuple ixs a))

-- {-# SPECIALISE Quick.sortInplace :: U.MVector s (OrdTuple '[Up 1] (Int, Word)) -> ST s () #-}

{-# SPECIALISE sortInplaceFM :: U.Unbox a => Sequential -> Median3or5 (OrdTuple '[Up 1] (a, Word)) -> U.MVector s (OrdTuple '[Up 1] (a, Word)) -> ST s () #-}

sortBySpec :: U.Unbox a => U.Vector (Strict.Pair a Word) -> U.Vector (Strict.Pair a Word)
sortBySpec = sortByOrdTuple (Proxy :: Proxy '[Up 1])

-- {-# SPECIALISE Quick.sortInplace :: U.MVector s (OrdTuple '[Up 1] (Int, Word)) -> ST s () #-}