Fix incorrect map behavior in devec (moonbitlang#541)

* Fix incorrect map behavior in devec * fix info * fix format
yj-qin · Jun 12, 2024 · dfc6f23 · dfc6f23
1 parent 1f8bcd4
commit dfc6f23
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 18 deletions.
diff --git a/devec/devec.mbt b/devec/devec.mbt
@@ -542,41 +542,53 @@ test "clear" {
 /// # Example
 /// ```
 /// let dv = Devec::[3, 4, 5]
-/// v.map(fn (x) {x + 1})
+/// dv.map(fn (x) {x + 1}) // Devec::[4, 5, 6]
 /// ```
-pub fn map[T](self : Devec[T], f : (T) -> T) -> Unit {
-  for i = 0; i < self.len; i = i + 1 {
-    self.buf[i] = f(self.buf[i])
+pub fn map[T, U](self : Devec[T], f : (T) -> U) -> Devec[U] {
+  if self.len == 0 {
+    Devec::[]
+  } else {
+    let buf : UninitializedArray[U] = UninitializedArray::make(self.len)
+    for i = 0; i < self.len; i = i + 1 {
+      buf[i] = f(self.buf[i])
+    }
+    Devec::{ buf, len: self.len, head: 0, tail: self.len - 1 }
   }
 }
 
 test "map" {
   let dv = Devec::[3, 4, 5]
-  dv.map(fn(x) { x + 1 })
-  @assertion.assert_eq(dv[0], 4)?
-  @assertion.assert_eq(dv[1], 5)?
-  @assertion.assert_eq(dv[2], 6)?
+  let dvp = dv.map(fn(x) { x + 1 })
+  @assertion.assert_eq(dvp[0], 4)?
+  @assertion.assert_eq(dvp[1], 5)?
+  @assertion.assert_eq(dvp[2], 6)?
 }
 
 /// Maps a function over the elements of the devec with index.
 ///
 /// # Example
 /// ```
 /// let dv = Devec::[3, 4, 5]
-/// v.mapi(fn (i, x) {x + i})
+/// dv.mapi(fn (i, x) {x + i}) // Devec::[3, 5, 7]
 /// ```
-pub fn mapi[T](self : Devec[T], f : (Int, T) -> T) -> Unit {
-  for i = 0; i < self.len; i = i + 1 {
-    self.buf[i] = f(i, self.buf[i])
+pub fn mapi[T, U](self : Devec[T], f : (Int, T) -> U) -> Devec[U] {
+  if self.len == 0 {
+    Devec::[]
+  } else {
+    let buf : UninitializedArray[U] = UninitializedArray::make(self.len)
+    for i = 0; i < self.len; i = i + 1 {
+      buf[i] = f(i, self.buf[i])
+    }
+    Devec::{ buf, len: self.len, head: 0, tail: self.len - 1 }
   }
 }
 
 test "mapi" {
   let dv = Devec::[3, 4, 5]
-  dv.mapi(fn(i, x) { x + i })
-  @assertion.assert_eq(dv[0], 3)?
-  @assertion.assert_eq(dv[1], 5)?
-  @assertion.assert_eq(dv[2], 7)?
+  let dvp = dv.mapi(fn(i, x) { x + i })
+  @assertion.assert_eq(dvp[0], 3)?
+  @assertion.assert_eq(dvp[1], 5)?
+  @assertion.assert_eq(dvp[2], 7)?
 }
 
 /// Test if the devec is empty.

diff --git a/devec/devec.mbti b/devec/devec.mbti
@@ -18,8 +18,8 @@ impl Devec {
   iter_revi[T](Self[T], (Int, T) -> Unit) -> Unit
   iteri[T](Self[T], (Int, T) -> Unit) -> Unit
   length[T](Self[T]) -> Int
-  map[T](Self[T], (T) -> T) -> Unit
-  mapi[T](Self[T], (Int, T) -> T) -> Unit
+  map[T, U](Self[T], (T) -> U) -> Self[U]
+  mapi[T, U](Self[T], (Int, T) -> U) -> Self[U]
   new[T]() -> Self[T]
   op_equal[T : Eq](Self[T], Self[T]) -> Bool
   op_get[T](Self[T], Int) -> T