Unique vectors

doc/spec/tour.kk.md

@@ -650,6 +650,47 @@ that are not possible for general mutable reference cells.
 [Read more about state and multiple resumptions &adown;][#sec-multi-resume]
 {.learn}
 
+### Local Mutable Vectors {#sec-vectors;}
+
+Koka provides a built-in `vector` type that is backed by a constant sized C array of boxed values.
+
+Vectors still provide an immutable interface, but are faster than lists for random access.
+
+When using vectors locally you can 'mutate' them similarly to local variables:
+
+```
+fun vectors()
+  var myvec := vector-init(10, fn(i) i)
+  myvec[5] := 42
+  myvec.map(show).join(", ").println
+```
+
+If there is another reference to the vector, the vector will be copied on assignment to form the new vector. The elements that were not updated are shared between the old and the new vector. In the following example, the `vec-copy` will be a vector that shares the same values as `myvec` except at location `5`.
+
+```
+fun vectors()
+  var myvec := vector-init(10, fn(i) i)
+  val vec-copy = myvec
+  myvec[5] := 42
+```
+Koka can even [reuse][#sec-fbip] the memory of the elements of the vector if there is only one reference to the element. However, when using the assignment syntax above, Koka does not release the element from the vector in time for the in-place reuse to occur. If you want to update an element with the opportunity for in-place reuse, you can use the update function:
+```
+fun reuse()
+  var myvec := vector-init(10, fn(i) [i - 1, i, i + 1])
+  std/core/types/@byref(myvec).update(0, fn(l) l.map(fn(x) x + 1))
+```
+
+Because the `myvec` is unique and the value at that position is also unique, the memory of both the old vector and the old value can be used in place thanks to Perceus reference counting.
+
+Vectors can also be used outside of local variables, but setting / updating the vector will return a new vector (reusing the old memory in place if the vector was unique).
+```
+fun reuse()
+  val myvec = vector-init(10, fn(i) [i - 1, i, i + 1])
+  val newvec = myvec.update(0, fn(l) l.map(fn(x) x + 1))
+  val newvec1 = newvec.set(1, [-1])
+  newvec1[0].println
+  newvec1[1].println
+```
 
 ### Reference Cells and Isolated state {#sec-runst}
 

lib/std/core/inline/vector.h

@@ -16,12 +16,11 @@
 kk_datatype_t  kk_vector_to_list(kk_vector_t v, kk_datatype_t tail, kk_context_t* ctx);
 kk_vector_t    kk_list_to_vector(kk_datatype_t xs, kk_context_t* ctx);
 
-static inline kk_unit_t kk_vector_unsafe_assign( kk_vector_t v, kk_ssize_t i, kk_box_t x, kk_context_t* ctx  ) {
+static inline kk_unit_t kk_vector_unsafe_assign_borrow( kk_vector_t v, kk_ssize_t i, kk_box_t x, kk_context_t* ctx  ) {
   kk_ssize_t len;
   kk_box_t* p = kk_vector_buf_borrow(v,&len,ctx);
   kk_assert(i < len);
   p[i] = x;
-  kk_vector_drop(v,ctx); // TODO: use borrowing
   return kk_Unit;
 }
 
@@ -30,3 +29,11 @@ static inline kk_box_t kk_vector_at_int_borrow( kk_vector_t v, kk_integer_t n, k
   kk_box_t b = kk_vector_at_borrow(v,kk_integer_clamp_ssize_t_borrow(n,ctx),ctx);
   return b;
 }
+
+static inline kk_box_t kk_vector_unsafe_extract_borrow(const kk_vector_t v, kk_ssize_t i, kk_context_t* ctx) {
+  kk_assert(i < kk_vector_len_borrow(v,ctx));
+  kk_box_t* p = kk_vector_buf_borrow(v, NULL, ctx);
+  kk_box_t res = p[i];
+  p[i] = kk_box_null();
+  return res;
+}

lib/std/core/vector.kk

@@ -29,29 +29,78 @@ inline extern unsafe-idx( ^v : vector<a>, index : ssize_t ) : total a
   cs inline "(#1)[#2]"
   js inline "(#1)[#2]"
 
-inline extern unsafe-assign : forall<a> ( v : vector<a>, i : ssize_t, x : a ) -> total ()
-  c "kk_vector_unsafe_assign"
+// Obtains the value transferring ownership and replaces the value with uninitialized value
+// Precondition: `v` is unique
+inline fip extern unsafe-extract( ^v : vector<a>, index : ssize_t ) : a
+  c "kk_vector_unsafe_extract_borrow"
+  cs inline "(#1)[#2]" // js / cs need to copy prior to calling this to satisfy the preconditions
+  js inline "(#1)[#2]"
+
+// Assign to an entry in a vector.
+inline fip extern unsafe-assign( ^v : vector<a>, i : ssize_t, x : a ): total ()
+  c "kk_vector_unsafe_assign_borrow"
   cs inline "(#1)[#2] = #3"
   js inline "(#1)[#2] = #3"
 
-inline extern unsafe-vector : forall<a> ( n : ssize_t ) -> total vector<a>
-  c  inline "kk_vector_alloc(#1,kk_box_null(),kk_context())"
-  cs inline "(new ##1[#1])"
-  js inline "Array(#1)"
-
+// Create a new vector of length `n`  with uninitialized elements.
 pub extern @unsafe-vector : forall<a> ( n : ssize_t ) -> total vector<a>
   c  inline "kk_vector_alloc(#1,kk_box_null(),kk_context())"
   cs inline "(new ##1[#1])"
   js inline "Array(#1)"
 
+// This is unsafe, since the behavior that depends on this check needs to be observationally equivalent
+inline fip extern unsafe-is-unique( ^v : vector<a> ) : bool
+  c inline "kk_datatype_is_unique(#1, kk_context())"
+  js inline "false"
+  cs inline "false"
+
+// Ensures the vector is unique, copying if necessary
+// This is technically not fip, but we pretend it to be since it will only copy when non-unique
+pub inline fip fun vec/ensure-unique(v: vector<a>): vector<a>
+  if v.unsafe-is-unique then v else v.unsafe-copy
+
+// Copies the vector
+// This is technically not fip, which is why we call it `unsafe-copy`.
+// However, we pretend it to be fip, since it is only called from `vec/ensure-unique` which will only copy if non-unique
+inline fip extern unsafe-copy(v: vector<a>): vector<a>
+  c "kk_vector_copy"
+  js inline "[...(#1)]"
+  cs inline "(#1).Clone()"
+
+// Copies the vector
+// If you only want to ensure uniqueness use `ensure-unique` instead, which will only copy if non-unique
+pub inline extern copy(v: vector<a>): vector<a>
+  c "kk_vector_copy"
+  js inline "[...(#1)]"
+  cs inline "(#1).Clone()"
+
 // Assign to an entry in a local `:vector` variable.
-pub inline extern assign/@index( ^self : local-var<s,vector<a>>, ^index : int, assigned : a ) : <local<s>,exn|e> ()
+// ```
+// var v := vector(10, 0) 
+// var v[0] := 1 // updates v in place
+// ```
+// TODO: Is this safe in cs / js? I think we might need to conservatively copy?
+pub inline extern assign/@index( ^self : local-var<h,vector<a>>, ^index : int, assigned : a ) : <local<h>> ()
   c  "kk_ref_vector_assign_borrow"
   cs inline "(#1)[(int)#2] = #3"  // check bounds?
   js inline "(#1).value[#2] = #3"       // todo: check bounds!
 
+// Assign to an entry in a local `:vector` variable.
+pub inline fun assign/update( ^self : local-var<h,vector<a>>, ^index : int, f : a -> <local<h>,div> a ) : <local<h>,div> ()
+  val i = index.ssize_t
+  // Future TODO: Somehow get rid of div & allow linear effects
+  ensure-unique(std/core/types/@byref(self))
+  val x = self.unsafe-extract(i)
+  val res = f(x)
+  self[index] := res
+
+// Ensures that a vector in a local variable is unique
+// Future TODO: Remove div (need hdiv constraint propagation)
+pub inline fun local/ensure-unique(^v: local-var<h,vector<a>>): <local<h>,div> ()
+  if v.unsafe-is-unique then () else v := unsafe-copy(v)
+
 // Length of a vector.
-inline extern lengthz( ^v : vector<a> ) : ssize_t
+inline fip extern lengthz( ^v : vector<a> ) : ssize_t
   c  "kk_vector_len_borrow"
   cs inline "((#1).Length)"
   js inline "((#1).length)"
@@ -74,13 +123,38 @@ pub inline extern unit/vector : forall<a> () -> vector<a>
   cs inline "new ##1[0]"
   js inline "[]"
 
-
 // Return the element at position `index`  in vector `v`.
 // Raise an out of bounds exception if `index < 0`  or `index >= v.length`.
 pub fun @index( ^v : vector<a>, ^index : int ) : exn a
   val idx = index.ssize_t()
   if idx < v.lengthz then unsafe-idx(v,idx) else throw("index out of bounds",ExnRange)
 
+// Sets the value at a particular location in a vector
+// If the index is out of bounds, an exception is thrown
+// If the vector is unique the value will be updated in place
+pub fun set(v: vector<a>, ^index: int, x: a): exn vector<a>
+  val idx = index.ssize_t()
+  if idx >= v.lengthz then 
+    throw("index out of bounds", ExnRange)
+  else
+    val v' = v.ensure-unique // Ensure unique prior to assigning
+    unsafe-assign(v', idx, x)
+    v'
+
+// Updates the value at a particular location in a vector
+// If the index is out of bounds, an exception is thrown
+// If the vector is unique the value will be updated in place
+pub fun update(v: vector<a>, ^index: int, ^f: a -> <exn|e> a): <exn|e> vector<a>
+  val idx = index.ssize_t()
+  if idx >= v.lengthz then 
+    throw("index out of bounds", ExnRange)
+  else
+    val v' = v.ensure-unique // Ensure unique before owning the value
+    val res = f(unsafe-extract(v',idx))
+    val v'' = v'.ensure-unique // Ensure unique after effects
+    unsafe-assign(v'', idx, res)
+    v''
+
 // Return the element at position `index` in vector `v`, or `Nothing` if out of bounds
 pub fun at( ^v : vector<a>, ^index : int ) : maybe<a>
   val idx = index.ssize_t
@@ -102,7 +176,7 @@ pub fun vector-init-total( ^n : int, f : int -> a ) : vector<a>
 // Create a new vector of length `n`  with initial elements given by function `f` which can have a control effect.
 pub fun vector-init( ^n : int, f : int -> e a ) : e vector<a>
   val len = n.ssize_t
-  val v = unsafe-vector(len)
+  val v = @unsafe-vector(len)
   forz( len ) fn(i)
     unsafe-assign(v,i,f(i.int))
   v
@@ -126,11 +200,43 @@ pub fun foreach-while( v : vector<a>, f : a -> e maybe<b> ) : e maybe<b>
     f(v.unsafe-idx(i))
 
 // Apply a function `f` to each element in a vector `v`
-pub fun map( v : vector<a>, f : a -> e b ) : e vector<b>
-  val w = unsafe-vector(v.length.ssize_t)
-  v.foreach-indexedz fn(i,x)
-    unsafe-assign(w,i,f(x))
-  w
+pub fip fun map( v : vector<a>, ^f : a -> e b ) : e vector<b>
+  map-rec(v.ensure-unique, f, 0.ssize_t).unsafe-vector-cast
+
+// Inner loop for `map` which does in place updates when unique, otherwise copies the vector
+// `f` is safe to have effects that do not resume (since the vector will be unique and freed)
+//      or to resume more than once (since the vector will be copied)
+//      the element type cast is safe for the same reasons
+// Precondition `v` is unique
+fip fun map-rec(v: vector<a>, ^f: a -> e b, i: ssize_t): e vector<b>
+  if i >= v.lengthz then
+    v.unsafe-vector-cast
+  else
+    // v is unique, so we can own the value at index i
+    val a = unsafe-extract(v, i)
+    // If the function `f` does not resume, the vector will be dropped
+    //  since it is guaranteed to be unique at this point,
+    //  it will be freed and the uninitialized value will not be observable 
+    val res = unsafe-cast(f(a))
+    // If the function `f` captures a resumption such that v is no longer unique, we need to copy the vector
+    val v' = v.ensure-unique
+    // We copy before assigning to avoid the result value being dupped during a non-unique copy
+    unsafe-assign(v', i, res)
+    map-rec(v'.pretend-decreasing, f, i.incr)
+
+// Only use in map-rec, which guarantees either:
+//  - all vector elements will eventually be converted to type `b` 
+//  - or the whole vector is unique and will be freed
+inline fip extern unsafe-cast( x : a ) : b
+  c inline "#1"
+  js inline "#1"
+  cs inline "#1"
+
+// Casts a vector type, this is unsafe and should only be used internally in `vector/map` 
+inline fip extern unsafe-vector-cast(v: vector<a>): vector<b>
+  c inline "#1"
+  js inline "#1"
+  cs inline "#1"
 
 // Convert a vector to a list.
 pub fun list( v : vector<a> ) : list<a>
@@ -174,8 +280,6 @@ fun for-whilez( n : ssize_t, action : (ssize_t) -> e maybe<a> ) : e maybe<a>
   rep(0.ssize_t)
 
 
-
-
 // Minimal set of operations that we need in `std/core`.
 inline fip extern ssize_t/(<=) : (ssize_t,ssize_t) -> bool
   inline "(#1 <= #2)"

test/cgen/vec-unique.kk

@@ -0,0 +1,129 @@
+// Update function
+fun update(xs: list<int>): console list<int>
+  println(if xs.unsafe-is-unique then "unique a" else "not unique a")
+  xs.map fn(x)
+    x + 1
+
+fun update-trace(xs: list<int>): list<int>
+  trace(if xs.unsafe-is-unique then "unique a" else "not unique a")
+  xs.map fn(x)
+    x + 1
+
+// This is unsafe, since the behavior that depends on this check needs to be observationally equivalent
+inline fip extern unsafe-is-unique( ^v : list<a> ) : bool
+  c inline "kk_datatype_is_unique(#1, kk_context())"
+
+effect choice
+  ctl choice() : int
+
+// https://github.com/koka-lang/koka/issues/544#issuecomment-2155184608 Anton's example
+fun antons-example()
+  // Nondeterministic choice, should copy vectors
+  var a := vector-init-total(3, fn(i) 0)
+  a[1] := 1
+  val xs =
+    with handler
+      return(x) [x]
+      ctl choice()
+        resume(0) ++ resume(1)
+    a.map(fn(x) abs(x - choice()))
+  xs.map(fn(v) v.map(show).join(",")).println
+  ()
+
+fun main()
+  println("Non-fip copy of vector and items")
+  val v = [[1], [2], [3], [4]].vector
+  val xs = v
+  // Copies the vector and the items are no longer unique
+  v.map(update).map(fn(x) x.show).join("").println
+  xs.map(fn(x) x.show).join("").println
+
+  // Items are unique, since the vector is unique
+  println("Fip updates to items in vector")
+  val v1 = [[1], [2], [3], [4]].vector
+  v1.map(update).map(fn(x) x.show).join("").println
+
+  println("Fip updates to vector, but only one items")
+  val v2 = [[1], [2], [3], [4]].vector
+  val v3 = v2.set(0, [-10])
+  val v4 = v3.update(1, fn(l) l.update())
+  v4.map(fn(x) x.show).join(",").println
+
+  println("Non-fip copies of vector in multiple resumptions")
+  antons-example()
+
+  // Fip update to item in local vector variable
+  println("Fip update to item in local vector variable")
+  var a1 := vector-init-total(3, fn(i) [i])
+  a1[1] := [5]
+  std/core/types/@byref(a1).update(2, update-trace)
+  a1.map(fn(x) x.show).join("").println
+
+  println("Non-fip update to item in local vector variable")
+  var a2 := vector-init-total(3, fn(i) [i])
+  val a3 = a2
+  std/core/types/@byref(a2).update(2, update-trace)
+  a2.map(fn(x) x.show).join("").println
+  a3.map(fn(x) x.show).join("").println
+
+
+// TODO: Needs to be fixed!!!
+extend type exception-info
+  ExnL(l: list<int>)
+
+fun update-throw(x: list<int>)
+  println(if x.unsafe-is-unique then "unique a" else "not unique a")
+  throw-exn(Exception("err", ExnL(x)))
+
+fun update-throw2(x: list<int>)
+  println(if x.unsafe-is-unique then "unique a" else "not unique a")
+  throw(x.show)
+
+// The nice thing about multiple references is that it will always be copied - so there is no issue
+// The problem with exceptions, is the reference count is too low
+// 
+// fun test-error1()
+//   println("Non-fip update to item in local vector variable")
+//   var a2 := vector-init-total(3, fn(i) [i])
+//   val a3 = a2
+//   try {
+//     update(std/core/types/@byref(a2), 2, update-throw)
+//     a2.map(fn(x) x.show).join("").println
+//   } fn(err) {
+//       match err.info 
+//         ExnL(l) -> println(if l.unsafe-is-unique then "unique a" else "not unique a")
+//       a2.foreach(fn(x) x.show.trace)
+//   }
+
+// fun test-error2()
+//   println("Non-fip update to item in local vector variable")
+//   val a4 = vector-init-total(3, fn(i) [i])
+//   var a5 := a4
+//   try {
+//     update(std/core/types/@byref(a5), 2, update-throw)
+//     a5.map(fn(x) x.show).join("").println
+//   } fn(err) {
+//       match err.info
+//         ExnL(l) -> println(if l.unsafe-is-unique then "unique a" else "not unique a")
+//       a5.foreach(fn(x) x.show.trace)
+//   }
+
+// fun test-error3()
+//   println("Non-fip update to item in local vector variable")
+//   val a6 = vector-init-total(3, fn(i) [i])
+//   var a7 := a6
+//   try {
+//     update(std/core/types/@byref(a7), 2, update-throw)
+//     a7.map(fn(x) x.show).join("").println
+//   } fn(err) {
+//       match err.info
+//         ExnL(l) -> println(if l.unsafe-is-unique then "unique a" else "not unique a")
+//       a7.foreach(fn(x) x.show.trace)
+//   }
+
+// This works for some reason?
+// fun test-error4()
+//   println("Non-fip update to item in local vector variable")
+//   val a6 = vector-init-total(3, fn(i) [i])
+//   var a7 := a6
+//   update(std/core/types/@byref(a7), 2, update-throw)