Transpose functions (#86)

rose-lang · Sep 11, 2023 · 1ee11a9 · 1ee11a9
1 parent 53267c2
commit 1ee11a9
Show file tree

Hide file tree

Showing 10 changed files with 1,155 additions and 22 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/interp/src/lib.rs b/crates/interp/src/lib.rs
@@ -16,7 +16,7 @@ pub enum Val {
     Bool(bool),
     F64(Cell<f64>),
     Fin(usize),
-    Ref(Rc<Val>),
+    Ref(Rc<Val>, Option<usize>),
     Array(Vals), // assume all indices are `Fin`
     Tuple(Vals),
 }
@@ -50,9 +50,35 @@ impl Val {
         }
     }
 
+    fn fin(&self) -> usize {
+        match self {
+            &Val::Fin(i) => i,
+            _ => unreachable!(),
+        }
+    }
+
+    fn get(&self, i: usize) -> &Self {
+        match self {
+            Val::Array(x) => &x[i],
+            Val::Tuple(x) => &x[i],
+            _ => unreachable!(),
+        }
+    }
+
+    fn slice(&self, i: usize) -> Self {
+        match self {
+            Val::Ref(x, None) => Val::Ref(Rc::clone(x), Some(i)),
+            Val::Ref(x, Some(j)) => Val::Ref(Rc::new(x.get(*j).clone()), Some(i)),
+            _ => unreachable!(),
+        }
+    }
+
     fn inner(&self) -> &Self {
         match self {
-            Val::Ref(x) => x.as_ref(),
+            Val::Ref(x, i) => match i {
+                None => x.as_ref(),
+                &Some(j) => x.get(j),
+            },
             _ => unreachable!(),
         }
     }
@@ -64,7 +90,7 @@ impl Val {
             &Self::Bool(x) => Self::Bool(x),
             Self::F64(_) => Self::F64(Cell::new(0.)),
             &Self::Fin(x) => Self::Fin(x),
-            Self::Ref(_) => unreachable!(),
+            Self::Ref(..) => unreachable!(),
             Self::Array(x) => Self::Array(collect_vals(x.iter().map(|x| x.zero()))),
             Self::Tuple(x) => Self::Tuple(collect_vals(x.iter().map(|x| x.zero()))),
         }
@@ -192,14 +218,8 @@ impl<'a, 'b, O: Opaque, T: Refs<'a, Opaque = O>> Interpreter<'a, 'b, O, T> {
                 _ => unreachable!(),
             },
 
-            &Expr::Slice { array, index } => match (self.get(array).inner(), self.get(index)) {
-                (Val::Array(v), &Val::Fin(i)) => v[i].clone(),
-                _ => unreachable!(),
-            },
-            &Expr::Field { tuple, member } => match self.get(tuple).inner() {
-                Val::Tuple(x) => x[member.member()].clone(),
-                _ => unreachable!(),
-            },
+            &Expr::Slice { array, index } => self.get(array).slice(self.get(index).fin()),
+            &Expr::Field { tuple, member } => self.get(tuple).slice(member.member()),
 
             &Expr::Unary { op, arg } => {
                 let x = self.get(arg);
@@ -257,8 +277,8 @@ impl<'a, 'b, O: Opaque, T: Refs<'a, Opaque = O>> Interpreter<'a, 'b, O, T> {
                 ))
             }
 
-            &Expr::Read { var } => Val::Ref(Rc::new(self.get(var).clone())),
-            &Expr::Accum { shape } => Val::Ref(Rc::new(self.get(shape).zero())),
+            &Expr::Read { var } => Val::Ref(Rc::new(self.get(var).clone()), None),
+            &Expr::Accum { shape } => Val::Ref(Rc::new(self.get(shape).zero()), None),
 
             &Expr::Ask { var } => self.get(var).inner().clone(),
             &Expr::Add { accum, addend } => {

diff --git a/crates/transpose/Cargo.toml b/crates/transpose/Cargo.toml
@@ -0,0 +1,8 @@
+[package]
+name = "rose-transpose"
+version = "0.0.0"
+edition = "2021"
+
+[dependencies]
+enumset = "1"
+rose = { path = "../core" }