{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE StrictData #-}

module Unison.Util.CyclicOrd where

import Data.Sequence qualified as S
import Data.Vector (Vector)
import Data.Vector qualified as V
import Unison.Prelude
import Unison.Util.CycleTable (CycleTable)
import Unison.Util.CycleTable qualified as CT

-- Same idea as `CyclicEq`, but for ordering.
class CyclicOrd a where
  -- Map from `Ref` ID to position in the stream
  -- If a ref is encountered again, we use its mapped ID
  cyclicOrd :: CycleTable Int Int -> CycleTable Int Int -> a -> a -> IO Ordering

bothOrd' ::
  (Ord a, CyclicOrd b) =>
  CT.CycleTable Int Int ->
  CT.CycleTable Int Int ->
  a ->
  a ->
  b ->
  b ->
  IO Ordering
bothOrd' :: forall a b.
(Ord a, CyclicOrd b) =>
CycleTable Int Int
-> CycleTable Int Int -> a -> a -> b -> b -> IO Ordering
bothOrd' CycleTable Int Int
h1 CycleTable Int Int
h2 a
a1 a
a2 b
b1 b
b2 = case a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
a1 a
a2 of
  Ordering
EQ -> CycleTable Int Int -> CycleTable Int Int -> b -> b -> IO Ordering
forall a.
CyclicOrd a =>
CycleTable Int Int -> CycleTable Int Int -> a -> a -> IO Ordering
cyclicOrd CycleTable Int Int
h1 CycleTable Int Int
h2 b
b1 b
b2
  Ordering
c -> Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ordering
c

bothOrd ::
  (CyclicOrd a, CyclicOrd b) =>
  CT.CycleTable Int Int ->
  CT.CycleTable Int Int ->
  a ->
  a ->
  b ->
  b ->
  IO Ordering
bothOrd :: forall a b.
(CyclicOrd a, CyclicOrd b) =>
CycleTable Int Int
-> CycleTable Int Int -> a -> a -> b -> b -> IO Ordering
bothOrd CycleTable Int Int
h1 CycleTable Int Int
h2 a
a1 a
a2 b
b1 b
b2 =
  CycleTable Int Int -> CycleTable Int Int -> a -> a -> IO Ordering
forall a.
CyclicOrd a =>
CycleTable Int Int -> CycleTable Int Int -> a -> a -> IO Ordering
cyclicOrd CycleTable Int Int
h1 CycleTable Int Int
h2 a
a1 a
a2 IO Ordering -> (Ordering -> IO Ordering) -> IO Ordering
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Ordering
b ->
    if Ordering
b Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ
      then CycleTable Int Int -> CycleTable Int Int -> b -> b -> IO Ordering
forall a.
CyclicOrd a =>
CycleTable Int Int -> CycleTable Int Int -> a -> a -> IO Ordering
cyclicOrd CycleTable Int Int
h1 CycleTable Int Int
h2 b
b1 b
b2
      else Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ordering
b

instance (CyclicOrd a) => CyclicOrd [a] where
  cyclicOrd :: CycleTable Int Int
-> CycleTable Int Int -> [a] -> [a] -> IO Ordering
cyclicOrd CycleTable Int Int
h1 CycleTable Int Int
h2 (a
x : [a]
xs) (a
y : [a]
ys) = CycleTable Int Int
-> CycleTable Int Int -> a -> a -> [a] -> [a] -> IO Ordering
forall a b.
(CyclicOrd a, CyclicOrd b) =>
CycleTable Int Int
-> CycleTable Int Int -> a -> a -> b -> b -> IO Ordering
bothOrd CycleTable Int Int
h1 CycleTable Int Int
h2 a
x a
y [a]
xs [a]
ys
  cyclicOrd CycleTable Int Int
_ CycleTable Int Int
_ [] [] = Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ordering
EQ
  cyclicOrd CycleTable Int Int
_ CycleTable Int Int
_ [] [a]
_ = Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ordering
LT
  cyclicOrd CycleTable Int Int
_ CycleTable Int Int
_ [a]
_ [] = Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ordering
GT

instance (CyclicOrd a) => CyclicOrd (S.Seq a) where
  cyclicOrd :: CycleTable Int Int
-> CycleTable Int Int -> Seq a -> Seq a -> IO Ordering
cyclicOrd CycleTable Int Int
h1 CycleTable Int Int
h2 Seq a
xs Seq a
ys = CycleTable Int Int
-> CycleTable Int Int -> [a] -> [a] -> IO Ordering
forall a.
CyclicOrd a =>
CycleTable Int Int -> CycleTable Int Int -> a -> a -> IO Ordering
cyclicOrd CycleTable Int Int
h1 CycleTable Int Int
h2 (Seq a -> [a]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq a
xs) (Seq a -> [a]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq a
ys)

instance (CyclicOrd a) => CyclicOrd (Vector a) where
  cyclicOrd :: CycleTable Int Int
-> CycleTable Int Int -> Vector a -> Vector a -> IO Ordering
cyclicOrd CycleTable Int Int
h1 CycleTable Int Int
h2 Vector a
xs Vector a
ys = Int
-> CycleTable Int Int
-> CycleTable Int Int
-> Vector a
-> Vector a
-> IO Ordering
forall {a}.
CyclicOrd a =>
Int
-> CycleTable Int Int
-> CycleTable Int Int
-> Vector a
-> Vector a
-> IO Ordering
go Int
0 CycleTable Int Int
h1 CycleTable Int Int
h2 Vector a
xs Vector a
ys
    where
      go :: Int
-> CycleTable Int Int
-> CycleTable Int Int
-> Vector a
-> Vector a
-> IO Ordering
go !Int
i !CycleTable Int Int
h1 !CycleTable Int Int
h2 !Vector a
xs !Vector a
ys =
        if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
xs Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
ys
          then Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ordering
EQ
          else
            if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
xs
              then Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ordering
LT
              else
                if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
ys
                  then Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ordering
GT
                  else do
                    Ordering
b <- CycleTable Int Int -> CycleTable Int Int -> a -> a -> IO Ordering
forall a.
CyclicOrd a =>
CycleTable Int Int -> CycleTable Int Int -> a -> a -> IO Ordering
cyclicOrd CycleTable Int Int
h1 CycleTable Int Int
h2 (Vector a
xs Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.! Int
i) (Vector a
ys Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.! Int
i)
                    if Ordering
b Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ
                      then Int
-> CycleTable Int Int
-> CycleTable Int Int
-> Vector a
-> Vector a
-> IO Ordering
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) CycleTable Int Int
h1 CycleTable Int Int
h2 Vector a
xs Vector a
ys
                      else Ordering -> IO Ordering
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Ordering
b