{-# language TemplateHaskellQuotes #-}
{-# language FlexibleContexts #-}
{-# language FlexibleInstances #-}
{-# language ScopedTypeVariables #-}
{-# language DataKinds #-}
{-# language TypeOperators #-}
{-# language EmptyCase #-}
{-# language DefaultSignatures #-}
{-# language BangPatterns #-}

module Language.Haskell.TH.TraverseCode
  ( TraverseCode (..)
  , sequenceCode
  , genericTraverseCode
  , genericSequenceCode
  ) where
import Generics.Linear
import Language.Haskell.TH.Syntax (Code, Lift (..), Exp (..), Quote, Name)
import qualified Language.Haskell.TH.Syntax as TH
import Language.Haskell.TH.Lib (conE)

-- for instances
import qualified Data.Functor.Product as FProd
import qualified Data.Functor.Sum as FSum
import Data.Functor.Identity
import qualified Data.Sequence.Internal as Seq
import Data.Coerce
import Control.Applicative
import qualified Data.Semigroup as S
import qualified Data.Monoid as M
import qualified Data.Complex as Complex
import qualified Data.Ord as Ord
import qualified Data.Functor.Compose as Compose
import qualified Data.Tree as Tree
import qualified Data.List.NonEmpty as NonEmpty
import GHC.Tuple (Solo)
import Data.Proxy (Proxy)
import qualified GHC.Generics as GHCGenerics
import qualified Data.Array as Ar
import qualified Data.Primitive.Array as PAr
import qualified Data.Primitive.SmallArray as PSmAr
import Data.Foldable (toList)

-- | Containers supporting \"traversal\" in 'Code'.
class TraverseCode t where
  -- | Given a container and a function to fill it with splices,
  -- produce a splice that will generate a container of their results.
  traverseCode :: Quote m => (a -> Code m b) -> t a -> Code m (t b)

  default traverseCode :: (Quote m, GTraverseCode (Rep1 t), Generic1 t) => (a -> Code m b) -> t a -> Code m (t b)
  traverseCode = (a -> Code m b) -> t a -> Code m (t b)
forall (m :: * -> *) (t :: * -> *) a b.
(Quote m, GTraverseCode (Rep1 t), Generic1 t) =>
(a -> Code m b) -> t a -> Code m (t b)
genericTraverseCode

-- | Given a container of splices, produce a splice that will generate a
-- container of their results.
sequenceCode :: (TraverseCode t, Quote m) => t (Code m a) -> Code m (t a)
sequenceCode :: forall (t :: * -> *) (m :: * -> *) a.
(TraverseCode t, Quote m) =>
t (Code m a) -> Code m (t a)
sequenceCode = (Code m a -> Code m a) -> t (Code m a) -> Code m (t a)
forall (t :: * -> *) (m :: * -> *) a b.
(TraverseCode t, Quote m) =>
(a -> Code m b) -> t a -> Code m (t b)
traverseCode Code m a -> Code m a
forall a. a -> a
id

-- | Like 'sequenceCode', but using the @"Generics.Linear".'Generic1'@ instance.
genericSequenceCode :: (Quote m, GTraverseCode (Rep1 t), Generic1 t) => t (Code m a) -> Code m (t a)
genericSequenceCode :: forall (m :: * -> *) (t :: * -> *) a.
(Quote m, GTraverseCode (Rep1 t), Generic1 t) =>
t (Code m a) -> Code m (t a)
genericSequenceCode = m Exp -> Code m (t a)
forall a (m :: * -> *). Quote m => m Exp -> Code m a
TH.unsafeCodeCoerce (m Exp -> Code m (t a))
-> (t (Code m a) -> m Exp) -> t (Code m a) -> Code m (t a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Code m a -> Code m a) -> Rep1 t (Code m a) -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCode f, Quote m) =>
(a -> Code m b) -> f a -> m Exp
gtraverseCode Code m a -> Code m a
forall a. a -> a
id (Rep1 t (Code m a) -> m Exp)
-> (t (Code m a) -> Rep1 t (Code m a)) -> t (Code m a) -> m Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t (Code m a) -> Rep1 t (Code m a)
forall {k} (f :: k -> *) (p :: k). Generic1 f => f p -> Rep1 f p
from1

-- | Like 'traverseCode', but using the @"Generics.Linear".'Generic1'@ instance.
genericTraverseCode :: (Quote m, GTraverseCode (Rep1 t), Generic1 t) => (a -> Code m b) -> t a -> Code m (t b)
genericTraverseCode :: forall (m :: * -> *) (t :: * -> *) a b.
(Quote m, GTraverseCode (Rep1 t), Generic1 t) =>
(a -> Code m b) -> t a -> Code m (t b)
genericTraverseCode a -> Code m b
f = m Exp -> Code m (t b)
forall a (m :: * -> *). Quote m => m Exp -> Code m a
TH.unsafeCodeCoerce (m Exp -> Code m (t b)) -> (t a -> m Exp) -> t a -> Code m (t b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Code m b) -> Rep1 t a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCode f, Quote m) =>
(a -> Code m b) -> f a -> m Exp
gtraverseCode a -> Code m b
f (Rep1 t a -> m Exp) -> (t a -> Rep1 t a) -> t a -> m Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t a -> Rep1 t a
forall {k} (f :: k -> *) (p :: k). Generic1 f => f p -> Rep1 f p
from1

class GTraverseCode f where
  gtraverseCode :: Quote m => (a -> Code m b) -> f a -> m Exp

instance (Datatype c, GTraverseCodeCon f) => GTraverseCode (D1 c f) where
  gtraverseCode :: forall (m :: * -> *) a b.
Quote m =>
(a -> Code m b) -> D1 c f a -> m Exp
gtraverseCode a -> Code m b
f m :: D1 c f a
m@(M1 f a
x) = String -> String -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeCon f, Quote m) =>
String -> String -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeCon String
pkg String
modl a -> Code m b
f f a
x
    where
      pkg :: String
pkg = D1 c f a -> String
forall {k} (d :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
       (f :: k1 -> *) (a :: k1).
Datatype d =>
t d f a -> String
packageName D1 c f a
m
      modl :: String
modl = D1 c f a -> String
forall {k} (d :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
       (f :: k1 -> *) (a :: k1).
Datatype d =>
t d f a -> String
moduleName D1 c f a
m

class GTraverseCodeCon f where
  gtraverseCodeCon :: Quote m => String -> String -> (a -> Code m b) -> f a -> m Exp

-- This instance seems totally useless, but it's obviously valid.
instance GTraverseCodeCon V1 where
  gtraverseCodeCon :: forall (m :: * -> *) a b.
Quote m =>
String -> String -> (a -> Code m b) -> V1 a -> m Exp
gtraverseCodeCon String
_pkg String
_modl a -> Code m b
_f V1 a
x = case V1 a
x of

instance (GTraverseCodeCon f, GTraverseCodeCon g) => GTraverseCodeCon (f :+: g) where
  gtraverseCodeCon :: forall (m :: * -> *) a b.
Quote m =>
String -> String -> (a -> Code m b) -> (:+:) f g a -> m Exp
gtraverseCodeCon String
pkg String
modl a -> Code m b
f (L1 f a
x) = String -> String -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeCon f, Quote m) =>
String -> String -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeCon String
pkg String
modl a -> Code m b
f f a
x
  gtraverseCodeCon String
pkg String
modl a -> Code m b
f (R1 g a
y) = String -> String -> (a -> Code m b) -> g a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeCon f, Quote m) =>
String -> String -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeCon String
pkg String
modl a -> Code m b
f g a
y

instance (Constructor c, GTraverseCodeFields f) => GTraverseCodeCon (C1 c f) where
  gtraverseCodeCon :: forall (m :: * -> *) a b.
Quote m =>
String -> String -> (a -> Code m b) -> C1 c f a -> m Exp
gtraverseCodeCon String
pkg String
modl a -> Code m b
f m :: C1 c f a
m@(M1 f a
x) = m Exp -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields (Name -> m Exp
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
conN) a -> Code m b
f f a
x
    where
      conBase :: String
conBase = C1 c f a -> String
forall {k} (c :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
       (f :: k1 -> *) (a :: k1).
Constructor c =>
t c f a -> String
conName C1 c f a
m
      conN :: Name
      conN :: Name
conN = String -> String -> String -> Name
TH.mkNameG_d String
pkg String
modl String
conBase

class GTraverseCodeFields f where
  gtraverseCodeFields :: Quote m => m Exp -> (a -> Code m b) -> f a -> m Exp

instance GTraverseCodeFields f => GTraverseCodeFields (S1 c f) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> S1 c f a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f (M1 f a
x) = m Exp -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f f a
x

instance (GTraverseCodeFields f, GTraverseCodeFields g) => GTraverseCodeFields (f :*: g) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> (:*:) f g a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f (f a
x :*: g a
y) =
    m Exp -> (a -> Code m b) -> g a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields (m Exp -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f f a
x) a -> Code m b
f g a
y

instance Lift p => GTraverseCodeFields (K1 i p) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> K1 i p a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (K1 p
x) = [| $c x |]

instance GTraverseCodeFields Par1 where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> Par1 a -> m Exp
gtraverseCodeFields m Exp
cc a -> Code m b
f (Par1 a
ca) = [| $cc $(TH.unTypeCode (f ca)) |]

instance GTraverseCodeFields U1 where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> U1 a -> m Exp
gtraverseCodeFields m Exp
cc a -> Code m b
_f U1 a
U1 = m Exp
cc

instance (GTraverseCodeFields f, TraverseCode g) => GTraverseCodeFields (f :.: g) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> (:.:) f g a -> m Exp
gtraverseCodeFields m Exp
cc a -> Code m b
f (Comp1 f (g a)
x) =
    m Exp -> (g a -> Code m (g b)) -> f (g a) -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields m Exp
cc ((a -> Code m b) -> g a -> Code m (g b)
forall (t :: * -> *) (m :: * -> *) a b.
(TraverseCode t, Quote m) =>
(a -> Code m b) -> t a -> Code m (t b)
traverseCode a -> Code m b
f) f (g a)
x

instance GTraverseCodeFields (URec Char) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> URec Char a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (UChar Char#
ch) = [| $c (UChar ch) |]

instance GTraverseCodeFields (URec Float) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> URec Float a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (UFloat Float#
ch) = [| $c (UFloat ch) |]

instance GTraverseCodeFields (URec Double) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> URec Double a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (UDouble Double#
ch) = [| $c (UDouble ch) |]

instance GTraverseCodeFields (URec Int) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> URec Int a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (UInt Int#
ch) = [| $c (UInt ch) |]

instance GTraverseCodeFields (URec Word) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> URec Word a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (UWord Word#
ch) = [| $c (UWord ch) |]

instance GTraverseCodeFields f => GTraverseCodeFields (MP1 m f) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> MP1 m f a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f (MP1 f a
x) = m Exp -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f f a
x

-- Would an instance for URec (Ptr ()) make any sense?

-- TraverseCode instances

instance TraverseCode Maybe
instance TraverseCode Identity
instance TraverseCode []
instance TH.Lift a => TraverseCode (Either a)

instance TH.Lift a => TraverseCode ((,) a)
instance (TH.Lift a, TH.Lift b) => TraverseCode ((,,) a b)
instance (TH.Lift a, TH.Lift b, TH.Lift c) => TraverseCode ((,,,) a b c)
instance (TH.Lift a, TH.Lift b, TH.Lift c, TH.Lift d) => TraverseCode ((,,,,) a b c d)
instance (TH.Lift a, TH.Lift b, TH.Lift c, TH.Lift d, TH.Lift e) => TraverseCode ((,,,,,) a b c d e)

instance (TraverseCode f, TraverseCode g) => TraverseCode (FProd.Product f g)
instance (TraverseCode f, TraverseCode g) => TraverseCode (FSum.Sum f g)
instance Lift a => TraverseCode (Const a)

instance TraverseCode V1
instance TraverseCode U1
instance (TraverseCode f, TraverseCode g) => TraverseCode (f :*: g)
instance (TraverseCode f, TraverseCode g) => TraverseCode (f :+: g)
instance TraverseCode f => TraverseCode (M1 i c f)
instance TraverseCode Par1
instance Lift a => TraverseCode (K1 i a)
instance (TraverseCode f, TraverseCode g) => TraverseCode (f :.: g)
instance (TraverseCode f, TraverseCode g) => TraverseCode (Compose.Compose f g)
instance TraverseCode f => TraverseCode (MP1 m f)
instance TraverseCode f => TraverseCode (GHCGenerics.Rec1 f)
instance (TraverseCode f, TraverseCode g) => TraverseCode (f GHCGenerics.:.: g)

-- The Elem instance isn't needed for the Seq instance
instance TraverseCode Seq.Elem
instance TraverseCode Seq.Digit
instance TraverseCode Seq.Node
instance TraverseCode Seq.FingerTree
instance TraverseCode Seq.Seq where
  -- This wonky way of doing it makes for a more compact
  -- splice.
  traverseCode :: forall (m :: * -> *) a b.
Quote m =>
(a -> Code m b) -> Seq a -> Code m (Seq b)
traverseCode a -> Code m b
f Seq a
s = [|| coerceFT $$(traverseCode f ft') ||]
    where
      ft' :: FingerTree a
ft' = Seq a -> FingerTree a
forall a. Seq a -> FingerTree a
coerceSeq ((() -> a -> a) -> Seq () -> Seq a -> Seq a
forall a b c. (a -> b -> c) -> Seq a -> Seq b -> Seq c
Seq.zipWith ((a -> () -> a) -> () -> a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> () -> a
forall a b. a -> b -> a
const) (Int -> () -> Seq ()
forall a. Int -> a -> Seq a
Seq.replicate (Seq a -> Int
forall a. Seq a -> Int
Seq.length Seq a
s) ()) Seq a
s)

coerceFT :: Seq.FingerTree a -> Seq.Seq a
coerceFT :: forall a. FingerTree a -> Seq a
coerceFT = FingerTree a -> Seq a
coerce
coerceSeq :: Seq.Seq a -> Seq.FingerTree a
coerceSeq :: forall a. Seq a -> FingerTree a
coerceSeq = Seq a -> FingerTree a
coerce

instance TraverseCode Seq.ViewL
instance TraverseCode Seq.ViewR

instance TraverseCode ZipList
instance TraverseCode Complex.Complex
instance TraverseCode S.First
instance TraverseCode M.First
instance TraverseCode S.Last
instance TraverseCode M.Last
instance TraverseCode S.Min
instance TraverseCode S.Max
instance TraverseCode Ord.Down
instance TraverseCode S.WrappedMonoid
instance TraverseCode S.Dual
instance TraverseCode S.Product
instance TraverseCode S.Sum
instance TraverseCode Solo
instance TraverseCode Tree.Tree
instance TraverseCode NonEmpty.NonEmpty
-- instance TraverseCode TH.TyVarBndr  (Lift Name isn't in base)
instance TraverseCode m => TraverseCode (WrappedMonad m)
instance TraverseCode (p a) => TraverseCode (WrappedArrow p a)
instance TH.Lift a => TraverseCode (S.Arg a)
instance TraverseCode Proxy
instance TraverseCode f => TraverseCode (M.Ap f)
instance TraverseCode f => TraverseCode (M.Alt f)
instance TraverseCode (URec Char)
instance TraverseCode (URec Double)
instance TraverseCode (URec Float)
instance TraverseCode (URec Int)
instance TraverseCode (URec Word)

instance (Ar.Ix i, TH.Lift i) => TraverseCode (Ar.Array i) where
  traverseCode :: forall (m :: * -> *) a b.
Quote m =>
(a -> Code m b) -> Array i a -> Code m (Array i b)
traverseCode a -> Code m b
f Array i a
xs = [|| Ar.listArray bnds $$csc ||]
    where
      csc :: Code m [b]
csc = (a -> Code m b) -> [a] -> Code m [b]
forall (t :: * -> *) (m :: * -> *) a b.
(TraverseCode t, Quote m) =>
(a -> Code m b) -> t a -> Code m (t b)
traverseCode a -> Code m b
f (Array i a -> [a]
forall i e. Array i e -> [e]
Ar.elems Array i a
xs)
      bnds :: (i, i)
bnds = Array i a -> (i, i)
forall i e. Array i e -> (i, i)
Ar.bounds Array i a
xs

instance TraverseCode PAr.Array where
  traverseCode :: forall (m :: * -> *) a b.
Quote m =>
(a -> Code m b) -> Array a -> Code m (Array b)
traverseCode a -> Code m b
f Array a
ary = case [a]
lst of
      [] -> [|| PAr.emptyArray ||]
      [a
x] -> [|| pure $$(f x) ||]
      a
x : [a]
xs -> [|| unsafeArrayFromListN len $$(f x) $$(traverseCode f xs) ||]
    where
      len :: Int
len = Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
ary
      lst :: [a]
lst = Array a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Array a
ary

instance TraverseCode PSmAr.SmallArray where
  traverseCode :: forall (m :: * -> *) a b.
Quote m =>
(a -> Code m b) -> SmallArray a -> Code m (SmallArray b)
traverseCode a -> Code m b
f SmallArray a
ary = case [a]
lst of
      [] -> [|| PSmAr.emptySmallArray ||]
      [a
x] -> [|| pure $$(f x) ||]
      a
x : [a]
xs -> [|| unsafeSmallArrayFromListN len $$(f x) $$(traverseCode f xs) ||]
    where
      len :: Int
len = SmallArray a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SmallArray a
ary
      lst :: [a]
lst = SmallArray a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList SmallArray a
ary

-- | Strictly create an array from a nonempty list (represented as
-- a first element and a list of the rest) of a known length. If the length
-- of the list does not match the given length, this makes demons fly
-- out of your nose.
unsafeArrayFromListN :: Int -> a -> [a] -> PAr.Array a
unsafeArrayFromListN :: forall a. Int -> a -> [a] -> Array a
unsafeArrayFromListN Int
n a
y [a]
ys =
  Int -> a -> (forall s. MutableArray s a -> ST s ()) -> Array a
forall a.
Int -> a -> (forall s. MutableArray s a -> ST s ()) -> Array a
PAr.createArray Int
n a
y ((forall s. MutableArray s a -> ST s ()) -> Array a)
-> (forall s. MutableArray s a -> ST s ()) -> Array a
forall a b. (a -> b) -> a -> b
$ \MutableArray s a
ma ->
    let go :: Int -> [a] -> ST s ()
go !Int
_ix [] = () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        go !Int
ix (a
x : [a]
xs) = do
            MutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
PAr.writeArray MutableArray s a
MutableArray (PrimState (ST s)) a
ma Int
ix a
x
            Int -> [a] -> ST s ()
go (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [a]
xs
    in Int -> [a] -> ST s ()
go Int
1 [a]
ys

-- | Strictly create an array from a nonempty list (represented as
-- a first element and a list of the rest) of a known length. If the length
-- of the list does not match the given length, this makes demons fly
-- out of your nose.
unsafeSmallArrayFromListN :: Int -> a -> [a] -> PSmAr.SmallArray a
unsafeSmallArrayFromListN :: forall a. Int -> a -> [a] -> SmallArray a
unsafeSmallArrayFromListN Int
n a
y [a]
ys =
  Int
-> a
-> (forall s. SmallMutableArray s a -> ST s ())
-> SmallArray a
forall a.
Int
-> a
-> (forall s. SmallMutableArray s a -> ST s ())
-> SmallArray a
PSmAr.createSmallArray Int
n a
y ((forall s. SmallMutableArray s a -> ST s ()) -> SmallArray a)
-> (forall s. SmallMutableArray s a -> ST s ()) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ \SmallMutableArray s a
ma ->
    let go :: Int -> [a] -> ST s ()
go !Int
_ix [] = () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        go !Int
ix (a
x : [a]
xs) = do
            SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
PSmAr.writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
ma Int
ix a
x
            Int -> [a] -> ST s ()
go (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [a]
xs
    in Int -> [a] -> ST s ()
go Int
1 [a]
ys