-- A fixed-size array.
data Vec (size ∷ Nat) a
   = Vec (SmallArray# a)
   | Pure a

-- Just `SmallArray#` in a Box so we can pass it around
data SA a = SA { unSA ∷ SmallArray# a }


-- A type-tagged integer that can either be computed at run-time, or
-- statically known if you have KnownNat
data KI (n ∷ Nat) = KI { unKI ∷ Int# }

-- like `Vec`, but also gives you the size as a type-tagged Int
pattern SVec s_val as <-
  Vec @s_ty as@(sizeofSmallArray#->KI @s_ty->s_val)

mkKI ∷ ∀ n -> KnownNat n => KI n
mkKI n | I# i <- fromIntegral $ natVal' (proxy# @n) = KI @n i

----

mkVec (KI @s sz) f = Vec @s $ -- Wrap the result immediately!
  -- If we did that inside the ST block, checking the constructor
  -- tag (Vec or Pure?) would force the array, preventing the
  -- more efficient `Pure` cases from actually saving any work.
  unSA $ runST $ ST \case
    s | (# s, sma #) <- newSmallArray# sz undefined s
      , (# s, sa  #) <- unsafeFreezeSmallArray# sma $
          forI 0# sz s \i -> writeSmallArray# sma i (f i)
     -> (# s, SA sa #)

forI start end s0 f = go start s0 where
  go i (s ∷ State# st)
    | 1# <- i <# end = go (i +# 1#) (f i s)
  go _ s = s

----

instance Functor (Vec s) where
  a <$ _ = Pure a
  fmap f = \case
    Pure a -> Pure (f a)
    SVec sa as -> mkVec sa \case
      i | (# a #) <- indexSmallArray# as i -> f a

-- The Applicative is zippy.
instance Applicative (Vec s) where
  pure = Pure; ma <* _ = ma; _ *> mb = mb
  liftA2 f = \cases
    (Pure a) bs -> fmap (a `f`) bs
    as (Pure b) -> fmap (`f` b) as
    (SVec sa as) (Vec bs) -> mkVec sa \case
      i | (# a #) <- indexSmallArray# as i
        , (# b #) <- indexSmallArray# bs i -> f a b

-- join picks the main diagonal [(0,0), (1,1), (2,2), ..]
instance Monad (Vec s) where
  m >>= f = case m of
    Pure a -> f a
    SVec sz sa -> mkVec sz \case
      i | (# a #) <- indexSmallArray# sa i ->
        case f a of
          Pure b -> b
          Vec sb | (# b #) <- indexSmallArray# sb i -> b

-- duplicate produces every rotation of the input vector:
--   transpose . toList . duplicate = toList . duplicate
instance Comonad (Vec s) where
  extract = \case
    Pure a -> a
    Vec sa | (# a #) <- indexSmallArray# sa 0# -> a
  extend f = \case
    Pure a -> Pure (f (Pure a))
    SVec sz@(KI s#) sa -> mkVec sz \i -> f $ mkVec sz \case
      j | (# a #) <- indexSmallArray# sa (remInt# (i +# j) s#) -> a

----

toFree = diagonal $ satisfy \a -> (a == (2,2)) <$ do
  lift $ vecFromListKI (mkKI 5) "lorem"
  tell $ Sum 1

-- ghci> print $ runWriterT $ runFree (toFree coords)
-- Vec [(Just (2,2),Sum {getSum = 15}),(Just (2,2),Sum {getSum = 15}),(Just (2,2),Sum {getSum = 15}),(Just (2,2),Sum {getSum = 15}),(Just (2,2),Sum {getSum = 15})]