Add some array builtins for more efficient array programming

parser-typechecker/src/Unison/Builtin.hs

@@ -255,7 +255,8 @@ builtinTypesSrc =
     B' "ClientSockAddr" CT.Data,
     B' "PinnedByteArray" CT.Data,
     B' "Integer" CT.Data,
-    B' "Natural" CT.Data
+    B' "Natural" CT.Data,
+    B' "UnboxedArray" CT.Data
   ]
 
 -- rename these to "builtin" later, when builtin means intrinsic as opposed to
@@ -676,29 +677,29 @@ builtinsSrc =
         --> nat
         --> Type.effect () [g, DD.exceptionType ()] unit,
     B "ImmutableArray.read" . forall1 "a" $ \a ->
-      iarrayt a --> nat --> Type.effect1 () (DD.exceptionType ()) a,
+      iarrayt a --> nat --> exnt a,
     B "ImmutableByteArray.read8" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "ImmutableByteArray.read16be" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "ImmutableByteArray.read24be" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "ImmutableByteArray.read32be" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "ImmutableByteArray.read40be" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "ImmutableByteArray.read64be" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "ImmutableByteArray.read16le" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "ImmutableByteArray.read24le" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "ImmutableByteArray.read32le" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "ImmutableByteArray.read40le" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "ImmutableByteArray.read64le" $
-      ibytearrayt --> nat --> Type.effect1 () (DD.exceptionType ()) nat,
+      ibytearrayt --> nat --> exnt nat,
     B "MutableArray.freeze!" . forall2 "g" "a" $ \g a ->
       marrayt g a --> Type.effect1 () g (iarrayt a),
     B "MutableByteArray.freeze!" . forall1 "g" $ \g ->
@@ -798,7 +799,76 @@ builtinsSrc =
     B "Natural.gteq" $ natural --> natural --> boolean,
     B "Natural.toFloat" $ natural --> float,
     B "Natural.isEven" $ natural --> boolean,
-    B "Natural.isOdd" $ natural --> boolean
+    B "Natural.isOdd" $ natural --> boolean,
+
+    B "ImmutableArray.at1s" $
+      forall2 "a" "b" $ \a b -> iarrayt (pair a b) --> iarrayt a,
+    B "ImmutableArray.at2s" $
+      forall3 "a" "b" "c" $ \a b c ->
+        iarrayt (pair a (pair b c)) --> iarrayt b,
+    B "ImmutableArray.chop" $
+      forall1 "a" $ \a ->
+        uarrayt nat --> uarrayt nat --> iarrayt a -->
+        exnt (iarrayt (iarrayt a)),
+    B "ImmutableArray.fromListAt1" $
+      forall2 "a" "b" $ \a b ->
+        list (pair a b) --> iarrayt a,
+    B "ImmutableArray.fromListAt2" $
+      forall3 "a" "b" "c" $ \a b c ->
+        list (pair a (pair b c)) --> iarrayt b,
+    B "ImmutableArray.intersectIx" $
+      forall1 "a" $ \a ->
+        iarrayt a --> iarrayt a --> tuple [uarrayt nat, uarrayt nat],
+    B "ImmutableArray.murmurHashesUntyped" $
+      forall1 "a" $ \a ->
+        iarrayt a --> uarrayt nat,
+    B "ImmutableArray.outerJoinIx" $
+      forall1 "a" $ \a ->
+        iarrayt a --> iarrayt a --> tuple [uarrayt nat, uarrayt nat],
+    B "ImmutableArray.pick" $
+      forall1 "a" $ \a ->
+        uarrayt nat --> iarrayt a --> exnt (iarrayt a),
+    B "ImmutableArray.pick1" $
+      forall1 "a" $ \a ->
+        uarrayt nat --> iarrayt a --> exnt (iarrayt a),
+    B "ImmutableArray.pick1Or" $
+      forall1 "a" $ \a ->
+        a --> uarrayt nat --> iarrayt a --> exnt (iarrayt a),
+    B "ImmutableArray.runsIx" $
+      forall1 "a" $ \a ->
+        iarrayt a --> tuple [uarrayt nat, uarrayt nat],
+    B "ImmutableArray.sortIx" $
+      forall1 "a" $ \a ->
+        iarrayt a --> uarrayt nat,
+    B "ImmutableArray.toLists" $
+      forall1 "a" $ \a ->
+        iarrayt (iarrayt a) --> iarrayt (list a),
+    B "ImmutableArray.toList" $
+      forall1 "a" $ \a ->
+        iarrayt a --> list a,
+    B "ImmutableArray.zipWithAppend" $
+      forall1 "a" $ \a ->
+        iarrayt (list a) --> iarrayt (list a) --> iarrayt (list a),
+    B "UnboxedArray.fromNatList" $
+      list nat --> uarrayt nat,
+    B "UnboxedArray.modR" $ uarrayt nat --> nat --> uarrayt nat,
+    B "UnboxedArray.multiplyR" $ uarrayt nat --> nat --> uarrayt nat,
+    B "UnboxedArray.divideR" $ uarrayt nat --> nat --> uarrayt nat,
+    B "UnboxedArray.occurrences" $
+      uarrayt nat --> iarrayt (uarrayt nat),
+    B "UnboxedArray.pick" $
+      forall1 "a" $ \a ->
+        uarrayt nat --> uarrayt a --> exnt (uarrayt a),
+    B "UnboxedArray.pick1" $
+      forall1 "a" $ \a ->
+        uarrayt nat --> uarrayt a --> exnt (uarrayt a),
+    B "UnboxedArray.pick1Or" $
+      forall1 "a" $ \a ->
+        a --> uarrayt nat --> uarrayt a --> exnt (uarrayt a),
+    B "UnboxedArray.size" $
+      forall1 "a" $ \a -> uarrayt a --> nat,
+    B "UnboxedArray.toList" $
+      forall1 "a" $ \a -> uarrayt a --> list a
   ]
     ++
     -- avoid name conflicts with Universal == < > <= >=
@@ -1118,6 +1188,17 @@ forall2 na nb body = Type.foralls () [a, b] (body ta tb)
     ta = Type.var () a
     tb = Type.var () b
 
+forall3 ::
+  Text -> Text -> Text -> (Type -> Type -> Type -> Type) -> Type
+forall3 na nb nc body = Type.foralls () [a,b,c] (body ta tb tc)
+  where
+    a = Var.named na
+    b = Var.named nb
+    c = Var.named nc
+    ta = Type.var () a
+    tb = Type.var () b
+    tc = Type.var () c
+
 forall4 ::
   Text ->
   Text ->
@@ -1191,6 +1272,12 @@ iarrayt a = Type.iarrayType () `app` a
 marrayt :: Type -> Type -> Type
 marrayt g a = Type.marrayType () `app` g `app` a
 
+uarrayt :: Type -> Type
+uarrayt a = Type.uarrayType () `app` a
+
+exnt :: Type -> Type
+exnt = Type.effect1 () (DD.exceptionType ())
+
 socket, threadId, handle, phandle, unit :: Type
 socket = Type.socket ()
 threadId = Type.threadId ()

parser-typechecker/src/Unison/KindInference/Generate.hs

@@ -377,6 +377,7 @@ builtinConstraintTree =
         (constrain (Type :-> Type))
         [ Type.list,
           Type.iarrayType,
+          Type.uarrayType,
           flip Type.ref Type.mvarRef,
           flip Type.ref Type.tvarRef,
           flip Type.ref Type.ticketRef,

unison-core/src/Unison/Type.hs

@@ -303,9 +303,10 @@ scopeRef, refRef :: TypeReference
 scopeRef = Reference.Builtin "Scope"
 refRef = Reference.Builtin "Ref"
 
-iarrayRef, marrayRef :: TypeReference
+iarrayRef, marrayRef, uarrayRef :: TypeReference
 iarrayRef = Reference.Builtin "ImmutableArray"
 marrayRef = Reference.Builtin "MutableArray"
+uarrayRef = Reference.Builtin "UnboxedArray"
 
 ibytearrayRef, mbytearrayRef :: TypeReference
 ibytearrayRef = Reference.Builtin "ImmutableByteArray"
@@ -425,6 +426,9 @@ ibytearrayType a = ref a ibytearrayRef
 mbytearrayType a = ref a mbytearrayRef
 pinnedByteArrayType a = ref a pinnedByteArrayRef
 
+uarrayType :: (Ord v) => a -> Type v a
+uarrayType a = ref a uarrayRef
+
 socket :: (Ord v) => a -> Type v a
 socket a = ref a socketRef
 

unison-runtime/package.yaml

@@ -104,6 +104,7 @@ library:
     - unliftio
     - unordered-containers
     - vector
+    - vector-algorithms
     - crypton-x509
     - crypton-x509-store
     - crypton-x509-system

unison-runtime/src/Unison/Runtime/ANF/POp.hs

@@ -149,6 +149,7 @@ data POp
   | CVLD -- validate
   | SDBX -- sandbox
   | VALU -- value
+  | VALS -- values
   | TLTT -- Term.Link.toText
   -- Debug
   | PRNT -- print
@@ -323,6 +324,7 @@ pOpCode op = case op of
   NOTB -> 145
   ANDB -> 146
   IORB -> 147
+  VALS -> 148
 
 pOpAssoc :: [(POp, Word16)]
 pOpAssoc = map (\op -> (op, pOpCode op)) [minBound .. maxBound]

unison-runtime/src/Unison/Runtime/Array.hs

@@ -26,6 +26,14 @@ module Unison.Runtime.Array
     readPrimArray,
     writePrimArray,
     indexPrimArray,
+    arrayMap,
+    arrayMapFromSeq,
+    arrayMapToPrim,
+    arrayToSeq,
+    SomePrimArr (..),
+    primArrayFromSeq,
+    primArrayMap,
+    primArrayToSeq,
   )
 where
 
@@ -55,7 +63,8 @@ import Data.Primitive.PrimArray as EPA hiding
   )
 import Data.Primitive.PrimArray qualified as PA
 import Data.Primitive.Types
-import Data.Word (Word8)
+import Data.Sequence (Seq, pattern (:<|), pattern Empty, (|>))
+import Data.Word (Word8, Word64)
 import GHC.IsList (toList)
 
 #ifdef ARRAY_CHECK
@@ -427,3 +436,88 @@ indexPrimArray = checkIPArray "indexPrimArray" PA.indexPrimArray
 
 byteArrayToList :: ByteArray -> [Word8]
 byteArrayToList = toList
+
+-- Builds a new array by applying a function to each element of an input
+-- array. The type is more like traverse, as it occurs in a `PrimMonad`
+-- for array creation, and we can permit the function to have effects in
+-- that monad. One thing this means is that you can throw IO exceptions if
+-- you want.
+arrayMap :: PrimMonad m => (a -> m b) -> PA.Array a -> m (Array b)
+arrayMap f ia = do
+  oa <- PA.newArray sz (error "arrayMap: dummy value")
+  let go n
+        | n < sz = do
+            PA.writeArray oa n =<< f (PA.indexArray ia n)
+            go (n+1)
+        | otherwise = PA.unsafeFreezeArray oa
+  go 0
+  where
+    sz = PA.sizeofArray ia
+{-# INLINE arrayMap #-}
+
+arrayMapFromSeq :: PrimMonad m => (a -> m b) -> Seq a -> m (Array b)
+arrayMapFromSeq f s = do
+  dst <- PA.newArray sz (error "arrayMapFromSeq: dummy value")
+  let go !_ Empty = PA.unsafeFreezeArray dst
+      go !n (x :<| xs) = do
+        PA.writeArray dst n =<< f x
+        go (n+1) xs
+  go 0 s
+  where
+    sz = length s
+{-# INLINE arrayMapFromSeq #-}
+
+arrayToSeq :: Array a -> Seq a
+arrayToSeq = fromList . toList
+
+-- TODO: more cases
+newtype SomePrimArr = NArr (PrimArray Word64)
+
+primArrayFromSeq :: (Prim a, PrimMonad m) => Seq a -> m (PrimArray a)
+primArrayFromSeq s = do
+  dst <- newPrimArray sz
+  let go !_ Empty = unsafeFreezePrimArray dst
+      go !n (x :<| xs) = writePrimArray dst n x *> go (n+1) xs
+  go 0 s
+  where
+    sz = length s
+{-# INLINE primArrayFromSeq #-}
+
+primArrayMap ::
+  (Prim a, Prim b, PrimMonad m) =>
+  (a -> m b) -> PrimArray a -> m (PrimArray b)
+primArrayMap f src = do
+  dst <- newPrimArray sz
+  let go n
+        | n >= sz = unsafeFreezePrimArray dst
+        | otherwise = do
+            writePrimArray dst n =<< f (indexPrimArray src n)
+            go (n+1)
+  go 0
+  where
+    sz = sizeofPrimArray src
+{-# INLINE primArrayMap #-}
+
+primArrayToSeq :: Prim a => PrimArray a -> Seq a
+primArrayToSeq src = go mempty 0
+  where
+    sz = sizeofPrimArray src
+
+    go acc n
+      | n < sz = go (acc |> indexPrimArray src n) (n+1)
+      | otherwise = acc
+{-# INLINE primArrayToSeq #-}
+
+arrayMapToPrim ::
+  (Prim b, PrimMonad m) => (a -> m b) -> Array a -> m (PrimArray b)
+arrayMapToPrim f src = do
+  dst <- newPrimArray sz
+  let go n
+        | n >= sz = unsafeFreezePrimArray dst
+        | otherwise = do
+            writePrimArray dst n =<< f (indexArray src n)
+            go (n+1)
+  go 0
+  where
+    sz = sizeofArray src
+{-# INLINE arrayMapToPrim #-}

unison-runtime/src/Unison/Runtime/Builtin.hs

@@ -643,6 +643,16 @@ murmur'hash instr =
   where
     (x, vl) = fresh
 
+-- custom code for hashing an entire array of values
+murmur'hashes :: ForeignOp
+murmur'hashes instr =
+  ([BX],)
+    . TAbss [x]
+    . TLetD vls BX (TPrm VALS [x])
+    $ TFOp instr [vls]
+  where
+    (x, vls) = fresh
+
 crypto'hmac :: ForeignOp
 crypto'hmac instr =
   ([BX, BX, BX],)
@@ -1227,6 +1237,8 @@ declareForeigns = do
 
   declareForeignWrap Untracked murmur'hash Universal_murmurHash
   declareForeignWrap Untracked murmur'hash Universal_murmurHashUntyped
+  declareForeignWrap Untracked murmur'hashes
+    ImmutableArray_murmurHashesUntyped
 
   declareForeign Tracked 1 IO_randomBytes
   declareForeign Untracked 1 Bytes_zlib_compress
@@ -1403,6 +1415,33 @@ declareForeigns = do
   declareForeign Untracked 1 Json_tryUnconsText
   declareForeign Untracked 3 Avro_decodeBinary
 
+  declareForeign Untracked 1 ImmutableArray_at1s
+  declareForeign Untracked 1 ImmutableArray_at2s
+  declareForeign Untracked 3 ImmutableArray_chop
+  declareForeign Untracked 1 ImmutableArray_fromList
+  declareForeign Untracked 1 ImmutableArray_fromListAt1
+  declareForeign Untracked 1 ImmutableArray_fromListAt2
+  declareForeign Untracked 2 ImmutableArray_intersectIx
+  declareForeign Untracked 2 ImmutableArray_outerJoinIx
+  declareForeign Untracked 2 ImmutableArray_pick
+  declareForeign Untracked 2 ImmutableArray_pick1
+  declareForeign Untracked 3 ImmutableArray_pick1Or
+  declareForeign Untracked 1 ImmutableArray_runsIx
+  declareForeign Untracked 1 ImmutableArray_toList
+  declareForeign Untracked 1 ImmutableArray_toLists
+  declareForeign Untracked 2 ImmutableArray_zipWithAppend
+  declareForeign Untracked 1 ImmutableArray_sortIx
+  declareForeign Untracked 1 UnboxedArray_fromNatList
+  declareForeign Untracked 2 UnboxedArray_modR
+  declareForeign Untracked 2 UnboxedArray_multiplyR
+  declareForeign Untracked 2 UnboxedArray_divideR
+  declareForeign Untracked 1 UnboxedArray_size
+  declareForeign Untracked 1 UnboxedArray_toList
+  declareForeign Untracked 1 UnboxedArray_occurrences
+  declareForeign Untracked 2 UnboxedArray_pick
+  declareForeign Untracked 2 UnboxedArray_pick1
+  declareForeign Untracked 3 UnboxedArray_pick1Or
+
 foreignDeclResults ::
   (Map ForeignFunc (Sandbox, SuperNormal Reference Symbol))
 foreignDeclResults =