From 85d6c2f07b88d6e97f08dc84e4beca4a0f77c377 Mon Sep 17 00:00:00 2001
From: Andrew Martin <andrew.thaddeus@gmail.com>
Date: Wed, 30 May 2018 13:29:09 -0400
Subject: [PATCH] Default prim implementations (#165)

* default implementations of setByteArray and setOffAddr along with examples.

* make types of default primitive methods more concise and add changelog entry

* add simple test for default implementations of set functions
---
 Data/Primitive/Types.hs | 48 ++++++++++++++++++++++++++++++++++++++++-
 cabal.project           |  1 -
 changelog.md            |  3 +++
 test/main.hs            | 21 ++++++++++++++++++
 4 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/Data/Primitive/Types.hs b/Data/Primitive/Types.hs
index 5fcdda5b..fd36ea0c 100644
--- a/Data/Primitive/Types.hs
+++ b/Data/Primitive/Types.hs
@@ -20,7 +20,7 @@
 
 module Data.Primitive.Types (
   Prim(..),
-  sizeOf, alignment,
+  sizeOf, alignment, defaultSetByteArray#, defaultSetOffAddr#,
 
   Addr(..),
   PrimStorable(..)
@@ -141,6 +141,52 @@ sizeOf x = I# (sizeOf# x)
 alignment :: Prim a => a -> Int
 alignment x = I# (alignment# x)
 
+-- | An implementation of 'setByteArray#' that calls 'writeByteArray#'
+-- to set each element. This is helpful when writing a 'Prim' instance
+-- for a multi-word data type for which there is no cpu-accelerated way
+-- to broadcast a value to contiguous memory. It is typically used
+-- alongside 'defaultSetOffAddr#'. For example:
+--
+-- > data Trip = Trip Int Int Int
+-- >
+-- > instance Prim Trip
+-- >   sizeOf# _ = 3# *# sizeOf# (undefined :: Int)
+-- >   alignment# _ = alignment# (undefined :: Int)
+-- >   indexByteArray# arr# i# = ...
+-- >   readByteArray# arr# i# = ...
+-- >   writeByteArray# arr# i# (Trip a b c) =
+-- >     \s0 -> case writeByteArray# arr# (3# *# i#) a s0 of
+-- >        s1 -> case writeByteArray# arr# ((3# *# i#) +# 1#) b s1 of
+-- >          s2 -> case writeByteArray# arr# ((3# *# i#) +# 2# ) c s2 of
+-- >            s3 -> s3
+-- >   setByteArray# = defaultSetByteArray#
+-- >   indexOffAddr# addr# i# = ...
+-- >   readOffAddr# addr# i# = ...
+-- >   writeOffAddr# addr# i# (Trip a b c) =
+-- >     \s0 -> case writeOffAddr# addr# (3# *# i#) a s0 of
+-- >        s1 -> case writeOffAddr# addr# ((3# *# i#) +# 1#) b s1 of
+-- >          s2 -> case writeOffAddr# addr# ((3# *# i#) +# 2# ) c s2 of
+-- >            s3 -> s3
+-- >   setOffAddr# = defaultSetOffAddr#
+defaultSetByteArray# :: Prim a => MutableByteArray# s -> Int# -> Int# -> a -> State# s -> State# s
+defaultSetByteArray# arr# i# len# ident = go 0#
+  where
+  go ix# s0 = if isTrue# (ix# <# len#)
+    then case writeByteArray# arr# (i# +# ix#) ident s0 of
+      s1 -> go (ix# +# 1#) s1
+    else s0
+
+-- | An implementation of 'setOffAddr#' that calls 'writeOffAddr#'
+-- to set each element. The documentation of 'defaultSetByteArray#'
+-- provides an example of how to use this.
+defaultSetOffAddr# :: Prim a => Addr# -> Int# -> Int# -> a -> State# s -> State# s
+defaultSetOffAddr# addr# i# len# ident = go 0#
+  where
+  go ix# s0 = if isTrue# (ix# <# len#)
+    then case writeOffAddr# addr# (i# +# ix#) ident s0 of
+      s1 -> go (ix# +# 1#) s1
+    else s0
+
 -- | Newtype that uses a 'Prim' instance to give rise to a 'Storable' instance.
 -- This type is intended to be used with the @DerivingVia@ extension available
 -- in GHC 8.6 and up. For example, consider a user-defined 'Prim' instance for
diff --git a/cabal.project b/cabal.project
index 24591f6f..6f92353a 100644
--- a/cabal.project
+++ b/cabal.project
@@ -3,4 +3,3 @@ packages: .
           ./bench
 package quickcheck-classes
   flags: -aeson -semigroupoids -semirings
-  
diff --git a/changelog.md b/changelog.md
index 57e59ced..53485f66 100644
--- a/changelog.md
+++ b/changelog.md
@@ -56,6 +56,9 @@
  * Add fold functions and map function to `Data.Primitive.UnliftedArray`.
    Add typeclass instances for `IsList`, `Ord`, and `Show`.
 
+ * Add `defaultSetByteArray#` and `defaultSetOffAddr#` to
+   `Data.Primitive.Types`.
+
 ## Changes in version 0.6.3.0
 
  * Add `PrimMonad` instances for `ContT`, `AccumT`, and `SelectT` from
diff --git a/test/main.hs b/test/main.hs
index 0c8b1f91..abec96df 100644
--- a/test/main.hs
+++ b/test/main.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE MagicHash #-}
 {-# LANGUAGE UnboxedTuples #-}
@@ -134,6 +135,9 @@ main = do
       , TQC.testProperty "foldlUnliftedArray'" (QCCL.foldlProp arrInt16 foldlUnliftedArray')
 #endif
       ]
+    , testGroup "DefaultSetMethod"
+      [ lawsToTest (QCC.primLaws (Proxy :: Proxy DefaultSetMethod))
+      ]
     -- , testGroup "PrimStorable"
     --   [ lawsToTest (QCC.storableLaws (Proxy :: Proxy Derived))
     --   ]
@@ -310,6 +314,23 @@ iforM_ xs0 f = go 0 xs0 where
   go !_ [] = return ()
   go !ix (x : xs) = f ix x >> go (ix + 1) xs
 
+newtype DefaultSetMethod = DefaultSetMethod Int16
+  deriving (Eq,Show,Arbitrary)
+
+instance Prim DefaultSetMethod where
+  sizeOf# _ = sizeOf# (undefined :: Int16)
+  alignment# _ = alignment# (undefined :: Int16)
+  indexByteArray# arr ix = DefaultSetMethod (indexByteArray# arr ix)
+  readByteArray# arr ix s0 = case readByteArray# arr ix s0 of
+    (# s1, n #) -> (# s1, DefaultSetMethod n #)
+  writeByteArray# arr ix (DefaultSetMethod n) s0 = writeByteArray# arr ix n s0
+  setByteArray# = defaultSetByteArray#
+  indexOffAddr# addr off = DefaultSetMethod (indexOffAddr# addr off)
+  readOffAddr# addr off s0 = case readOffAddr# addr off s0 of
+    (# s1, n #) -> (# s1, DefaultSetMethod n #)
+  writeOffAddr# addr off (DefaultSetMethod n) s0 = writeOffAddr# addr off n s0
+  setOffAddr# = defaultSetOffAddr#
+
 -- TODO: Uncomment this out when GHC 8.6 is release. Also, uncomment
 -- the corresponding PrimStorable test group above.
 --