Merge pull request #361 from koto-lang/koto-object-index-mut

Mutable indexing improvements

Author: irh

CHANGELOG.md

@@ -17,6 +17,7 @@ The Koto project adheres to
 - `export` can be used with multi-assignment expressions.
   - E.g. expressions like `export a, b, c = foo()` are now allowed.
 - Maps now support `[]` indexing, returning the Nth entry as a tuple.
+  - Entries can also be replaced by index by assigning a key/value tuple.
 - Objects that implement `KotoObject::call` can now be used in operations that
   expect functions.
   - `KotoObject::is_callable` has been added to support this, and needs to be
@@ -33,6 +34,8 @@ The Koto project adheres to
   - The macros in `koto_derive` have been updated to support generics.
     - `KotoField` has been added to reduce boilerplate when using the derive
       macros.
+- `KotoObject::index_mut` has been added to allow objects to support mutable
+  indexing operations.
 - `TryFrom<KValue>` has been implemented for some `core` and `std` types,
   including `bool`, string, and number types.
 

crates/bytecode/src/compiler.rs

@@ -2680,7 +2680,7 @@ impl Compiler {
                 }
                 ChainNode::Index(_) => {
                     let index = index.unwrap(self)?;
-                    self.push_op(SetIndex, &[container_register, index, value_register]);
+                    self.push_op(IndexMut, &[container_register, index, value_register]);
                 }
                 _ => {}
             }

crates/bytecode/src/instruction.rs

@@ -270,7 +270,7 @@ pub enum Instruction {
         value: u8,
         index: u8,
     },
-    SetIndex {
+    IndexMut {
         register: u8,
         index: u8,
         value: u8,
@@ -693,13 +693,13 @@ impl fmt::Debug for Instruction {
                 f,
                 "Index\t\tresult: {register}\tvalue: {value}\tindex: {index}"
             ),
-            SetIndex {
+            IndexMut {
                 register,
                 index,
                 value,
             } => write!(
                 f,
-                "SetIndex\tregister: {register}\tindex: {index}\tvalue: {value}"
+                "IndexMut\tregister: {register}\tindex: {index}\tvalue: {value}"
             ),
             MapInsert {
                 register,

crates/bytecode/src/instruction_reader.rs

@@ -386,7 +386,7 @@ impl Iterator for InstructionReader {
                 value: get_u8!(),
                 index: get_u8!(),
             }),
-            Op::SetIndex => Some(SetIndex {
+            Op::IndexMut => Some(IndexMut {
                 register: get_u8!(),
                 index: get_u8!(),
                 value: get_u8!(),

crates/bytecode/src/op.rs

@@ -413,7 +413,7 @@ pub enum Op {
     /// Sets a contained value via index
     ///
     /// `[*indexable, *value, *index]`
-    SetIndex,
+    IndexMut,
 
     /// Inserts a key/value entry into a map
     ///

crates/cli/docs/language_guide.md

@@ -643,6 +643,15 @@ print! m[1]
 check! ('oranges', 99)
 ```
 
+Entries can also be replaced by assigning a key/value tuple to the entry's index.
+
+```koto
+m = {apples: 42, oranges: 99, lemons: 63}
+m[1] = ('pears', 123)
+print! m
+check! {apples: 42, pears: 123, lemons: 63}
+```
+
 ### Shorthand Values
 
 Koto supports a shorthand notation when creating maps with inline syntax. 

crates/cli/docs/libs/color.md

@@ -214,15 +214,18 @@ check! Color(RGB, r: 0.2, g: 0.4, b: 0.3, a: 0.5)
 
 The `color` module's core color type.
 
+An `alpha` value is always present as the color's fourth component.
+
 The color may belong to various different color spaces, 
 with the space's components available via iteration or indexing.
 
-An `alpha` value is always present as the fourth component.
+Components can be modified via index.
+
 
 ### Example
 
 ```koto
-r, g, b = color('yellow')
+r, g, b = color 'yellow'
 print! r, g, b
 check! (1.0, 1.0, 0.0)
 
@@ -232,6 +235,13 @@ check! (90.0, 0.5, 0.25, 1.0)
 
 print! color('red')[0]
 check! 1.0
+
+print! c = color.oklch 0.5, 0.1, 180
+check! Color(Oklch, l: 0.5, c: 0.1, h: 180, a: 1)
+c[0] = 0.25 # Set the lightness component to 0.25
+c[3] = 0.1 # Set the alpha component to 0.1
+print c
+check! Color(Oklch, l: 0.25, c: 0.1, h: 180, a: 0.1)
 ```
 
 ## Color.mix

crates/runtime/src/types/object.rs

@@ -116,7 +116,7 @@ pub trait KotoObject: KotoType + KotoCopy + KotoEntries + KotoSend + KotoSync +
     ///
     /// By default, the object's type is used as the display string.
     ///
-    /// The [`DisplayContext`] is used to append strings to the result, and also provides context
+    /// The [`DisplayContext`] is used to append strings to the result, and provides information
     /// about any parent containers.
     fn display(&self, ctx: &mut DisplayContext) -> Result<()> {
         ctx.append(self.type_string());
@@ -130,6 +130,13 @@ pub trait KotoObject: KotoType + KotoCopy + KotoEntries + KotoSend + KotoSync +
         unimplemented_error("@index", self.type_string())
     }
 
+    /// Called when mutating an object via indexing, e.g. `x[0] = 99`
+    ///
+    /// See also: [KotoObject::size]
+    fn index_mut(&mut self, _index: &KValue, _value: &KValue) -> Result<()> {
+        unimplemented_error("@index_mut", self.type_string())
+    }
+
     /// Called when checking for the number of elements contained in the object
     ///
     /// The size should represent the maximum valid index that can be passed to

crates/runtime/src/vm.rs

@@ -892,11 +892,11 @@ impl KotoVm {
                 value,
                 index,
             } => self.run_index(register, value, index)?,
-            SetIndex {
+            IndexMut {
                 register,
                 index,
                 value,
-            } => self.run_set_index(register, index, value)?,
+            } => self.run_index_mut(register, index, value)?,
             MapInsert {
                 register,
                 key,
@@ -2221,7 +2221,7 @@ impl KotoVm {
         import_result
     }
 
-    fn run_set_index(
+    fn run_index_mut(
         &mut self,
         indexable_register: u8,
         index_register: u8,
@@ -2230,8 +2230,8 @@ impl KotoVm {
         use KValue::*;
 
         let indexable = self.clone_register(indexable_register);
-        let index_value = self.clone_register(index_register);
-        let value = self.clone_register(value_register);
+        let index_value = self.get_register(index_register);
+        let value = self.get_register(value_register);
 
         match indexable {
             List(list) => {
@@ -2240,24 +2240,50 @@ impl KotoVm {
                 match index_value {
                     Number(index) => {
                         let u_index = usize::from(index);
-                        if index >= 0.0 && u_index < list_len {
-                            list_data[u_index] = value;
+                        if *index >= 0.0 && u_index < list_len {
+                            list_data[u_index] = value.clone();
                         } else {
-                            return runtime_error!("Index '{index}' not in List");
+                            return runtime_error!("Invalid index ({index})");
                         }
                     }
                     Range(range) => {
                         for i in range.indices(list_len) {
                             list_data[i] = value.clone();
                         }
                     }
-                    unexpected => return unexpected_type("index", &unexpected),
+                    unexpected => return unexpected_type("Number or Range", unexpected),
                 }
+                Ok(())
             }
-            unexpected => return unexpected_type("a mutable indexable value", &unexpected),
-        };
-
-        Ok(())
+            Map(map) => match index_value {
+                Number(index) => {
+                    let mut map_data = map.data_mut();
+                    let map_len = map_data.len();
+                    let u_index = usize::from(index);
+                    if *index >= 0.0 && u_index < map_len {
+                        match value {
+                            Tuple(new_entry) if new_entry.len() == 2 => {
+                                let key = ValueKey::try_from(new_entry[0].clone())?;
+                                // There's no API on IndexMap for replacing an entry,
+                                // so use swap_remove_index to remove the old entry,
+                                // then insert the new entry at the end of the map,
+                                // followed by swap_indices to swap the new entry back into position.
+                                map_data.swap_remove_index(u_index);
+                                map_data.insert(key, new_entry[1].clone());
+                                map_data.swap_indices(u_index, map_len - 1);
+                                Ok(())
+                            }
+                            unexpected => unexpected_type("Tuple with 2 elements", unexpected),
+                        }
+                    } else {
+                        runtime_error!("Invalid index ({index})")
+                    }
+                }
+                unexpected => unexpected_type("Number", unexpected),
+            },
+            Object(o) => o.try_borrow_mut()?.index_mut(index_value, value),
+            unexpected => unexpected_type("a mutable indexable value", &unexpected),
+        }
     }
 
     fn validate_index(&self, n: KNumber, size: Option<usize>) -> Result<usize> {

crates/runtime/tests/object_tests.rs

@@ -140,6 +140,22 @@ mod objects {
             }
         }
 
+        fn index_mut(&mut self, index: &KValue, value: &KValue) -> Result<()> {
+            match index {
+                KValue::Number(index) => {
+                    assert_eq!(usize::from(index), 0);
+                    match value {
+                        KValue::Number(value) => {
+                            self.x = value.into();
+                            Ok(())
+                        }
+                        unexpected => unexpected_type("Number as value", unexpected),
+                    }
+                }
+                unexpected => unexpected_type("Number as index", unexpected),
+            }
+        }
+
         fn size(&self) -> Option<usize> {
             Some(self.x.unsigned_abs() as usize)
         }
@@ -749,6 +765,24 @@ match make_object 10
 ";
             test_object_script(script, 45);
         }
+
+        #[test]
+        fn index_mut_assign() {
+            let script = "
+x = make_object 100
+x[0] = 23
+";
+            test_object_script(script, 23);
+        }
+
+        #[test]
+        fn index_mut_compound_assign() {
+            let script = "
+x = make_object 100
+x[0] += 1
+";
+            test_object_script(script, 101);
+        }
     }
 
     #[test]

koto/tests/libs/color.koto

@@ -1,12 +1,11 @@
 import color
 
-assert_color_near = |a, b|
+assert_color_near = |(a0, a1, a2, a3), (b0, b1, b2, b3)|
   allowed_error = 1.0e-3
-  for component_a, component_b in a.zip b
-    assert_near component_a, component_b, allowed_error
-    assert_near component_a, component_b, allowed_error
-    assert_near component_a, component_b, allowed_error
-    assert_near component_a, component_b, allowed_error
+  assert_near a0, b0, allowed_error
+  assert_near a1, b1, allowed_error
+  assert_near a2, b2, allowed_error
+  assert_near a3, b3, allowed_error
 
 @tests =
   @test rgb: ||
@@ -38,6 +37,11 @@ assert_color_near = |a, b|
     assert_eq (color 'blue')[1], 0
     assert_eq (color 'blue')[2], 1
 
+  @test index_mut: ||
+    c = color 'red'
+    c[0] = 0.5
+    assert_eq c[0], 0.5
+
   @test iterator: ||
     assert_eq (color 'blue').to_tuple(), (0, 0, 1, 1)
 

libs/color/src/color.rs

@@ -32,10 +32,10 @@ impl Color {
         s.parse::<palette::Srgb<u8>>().ok().map(Self::from)
     }
 
-    pub fn get_component(&self, n: usize) -> Option<f32> {
+    pub fn get_component(&self, index: usize) -> Option<f32> {
         use Color::*;
 
-        let result = match (self, n) {
+        let result = match (self, index) {
             (Srgb(c), 0) => c.color.red,
             (Srgb(c), 1) => c.color.green,
             (Srgb(c), 2) => c.color.blue,
@@ -62,6 +62,36 @@ impl Color {
         Some(result)
     }
 
+    pub fn set_component(&mut self, index: usize, value: f32) -> Result<()> {
+        use Color::*;
+
+        match (self, index) {
+            (Srgb(c), 0) => c.color.red = value,
+            (Srgb(c), 1) => c.color.green = value,
+            (Srgb(c), 2) => c.color.blue = value,
+            (Srgb(c), 3) => c.alpha = value,
+            (Hsl(c), 0) => c.color.hue = value.into(),
+            (Hsl(c), 1) => c.color.saturation = value,
+            (Hsl(c), 2) => c.color.lightness = value,
+            (Hsl(c), 3) => c.alpha = value,
+            (Hsv(c), 0) => c.color.hue = value.into(),
+            (Hsv(c), 1) => c.color.saturation = value,
+            (Hsv(c), 2) => c.color.value = value,
+            (Hsv(c), 3) => c.alpha = value,
+            (Oklab(c), 0) => c.color.l = value,
+            (Oklab(c), 1) => c.color.a = value,
+            (Oklab(c), 2) => c.color.b = value,
+            (Oklab(c), 3) => c.alpha = value,
+            (Oklch(c), 0) => c.color.l = value,
+            (Oklch(c), 1) => c.color.chroma = value,
+            (Oklch(c), 2) => c.color.hue = value.into(),
+            (Oklch(c), 3) => c.alpha = value,
+            _ => return runtime_error!("Invalid component index ({index})"),
+        }
+
+        Ok(())
+    }
+
     pub fn color_space_str(&self) -> &str {
         use Color::*;
 
@@ -174,6 +204,20 @@ impl KotoObject for Color {
         }
     }
 
+    fn size(&self) -> Option<usize> {
+        // All current color spaces have 4 components
+        Some(4)
+    }
+
+    fn index_mut(&mut self, index: &KValue, value: &KValue) -> Result<()> {
+        use KValue::Number;
+
+        match (index, value) {
+            (Number(index), Number(value)) => self.set_component(index.into(), value.into()),
+            _ => unexpected_args("two Numbers", &[index.clone(), value.clone()]),
+        }
+    }
+
     fn is_iterable(&self) -> IsIterable {
         IsIterable::Iterable
     }