{-# LANGUAGE GHC2024 #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE UnicodeSyntax #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE ImpredicativeTypes #-}
-- base
import GHC.TypeLits
import Prelude (type (~), (==), Int, Maybe (Just, Nothing))
-- constraint
import Data.Constraint.Nat    (minusNat, plusMinusInverse1, plusNat)
import Data.Constraint        (HasDict, (\\), Dict (Dict))
import Data.Constraint.Unsafe (unsafeAxiom)
-- linear-base
import Prelude.Linear qualified as L
import Data.V.Linear  qualified as VL
import Unsafe.Linear  qualified as UnsafeLinear


-- Linear version of (\\)
(\\%) :: HasDict c e => (c => r) ⊸ e ⊸ r
(\\%) = UnsafeLinear.toLinear2 (\\)
infixl 1 \\%

case_non_zero_nat :: forall n a r p.
  KnownNat n =>
  SNat n ->
  a %p ->
  (                    a %p -> Dict (n  ~ 0) -> r) ->
  (KnownNat (n - 1) => a %p -> Dict (1 <= n) -> r)
  -> r
case_non_zero_nat sn a f1 f2 =
  if fromSNat sn == 0
  then let dict = unsafeAxiom @(n ~ 0) in f1 a dict
  else let dict = unsafeAxiom @(1 <= n) in f2 a dict \\% minusNat @n @1 \\% dict

f :: forall n. KnownNat n => VL.V n Int %1 -> (Maybe Int, VL.V n Int)
f vec_ = case_non_zero_nat (SNat @n) vec_
  (\vec _ -> (Nothing, vec))
  (\vec ev -> let !(a, vec') = VL.uncons vec \\% ev
                  !(L.Ur a') = L.move a
              in (Just a', VL.cons a' vec'))

main = do
  L.print (f (VL.make 4 2))
  L.print (f (VL.make))