Record Updates Allow Arbitrary Expressions, Update Pollen Tests (#141)

Co-authored-by: Susan Garry <[email protected]> Co-authored-by: Susan Garry <[email protected]> Co-authored-by: anshumanmohan <[email protected]> Co-authored-by: Adrian Sampson <[email protected]>
cucapra · Sep 28, 2023 · 4825794 · 4825794
1 parent 6fc0854
commit 4825794
Show file tree

Hide file tree

Showing 12 changed files with 830 additions and 15 deletions.
diff --git a/pollen/src/ast/mod.rs b/pollen/src/ast/mod.rs
@@ -75,12 +75,16 @@ pub enum Expr {
         op: UOp,
         expr: Box<Expr>
     },
+    ArrayAccess {
+        expr: Box<Expr>,
+        idx: Box<Expr>
+    },
     Record {
         typ: Typ,
         fields: Vec<RecordField>
     },
     RecordUpdate {
-        parent: Id,
+        parent: Box<Expr>,
         fields: Vec<RecordField>
     },
     Tuple {

diff --git a/pollen/src/main.rs b/pollen/src/main.rs
@@ -25,6 +25,7 @@ lazy_static! {
 
         // Precedence is defined lowest to highest
         PrattParser::new()
+            .op(Op::postfix(Rule::array_access))
             .op(Op::infix(Rule::or, Left))
             .op(Op::infix(Rule::and, Left))
             .op(Op::infix(Rule::eq, Left) | Op::infix(Rule::neq, Left))
@@ -544,13 +545,13 @@ fn parse_expr(expression: Pairs<Rule>) -> Expr {
                 }
             },
             Rule::record_update_lit => {
-                // record_update_lit looks like { r1 with f1: e1, ..., fn:en }
+                // record_update_lit looks like { e1 with f1: e1, ..., fn:en }
                 let mut inner = primary.into_inner();
                 let parent = {
                     let Some(pair) = inner.next() else {
-                        unreachable!("An if statement requires a guard")
+                        unreachable!("Record update requires a record")
                     };
-                    parse_id(pair)
+                    parse_expr(pair.into_inner())
                 };
                 let mut fields = Vec::new();
                 while let Some(pair) = inner.next() {
@@ -569,7 +570,7 @@ fn parse_expr(expression: Pairs<Rule>) -> Expr {
                     );
                 }
                 Expr::RecordUpdate {
-                    parent: parent,
+                    parent: Box::new(parent),
                     fields: fields
                 }
             },
@@ -625,10 +626,22 @@ fn parse_expr(expression: Pairs<Rule>) -> Expr {
                 expr: Box::new(exp),
             }
         })
-        .map_postfix(|lhs, op| {
-            match op.as_rule() {
-                rule => unreachable!("{:?} not recognized as a postfix", rule),
-            }
+        .map_postfix(|exp, op| {
+            let idx_expr = match op.as_rule() {
+                Rule::array_access => {
+                    let mut inner = op.into_inner();
+                    let Some(expr) = inner.next() else {
+                        unreachable!("array access requires an expression")
+                    };
+                    parse_expr(expr.into_inner())
+                }
+                rule => unreachable!("{:?} not recognized as a postfix", rule)
+            };
+
+                Expr::ArrayAccess {
+                    expr: Box::new(exp),
+                    idx: Box::new(idx_expr)
+                }
         })
         .map_infix(|lhs, op, rhs| {
             enum OpType {

diff --git a/pollen/src/pollen.pest b/pollen/src/pollen.pest
@@ -73,7 +73,7 @@ record_lit = { typ ~ "{" ~
                 identifier ~ ":" ~ expr  ~ 
                 ("," ~ identifier ~ ":" ~ expr)* ~ "}"
              }
-record_update_lit = { "{" ~ identifier ~ "with" ~
+record_update_lit = { "{" ~ expr ~ "with" ~
                 identifier ~ ":" ~ expr  ~ 
                 ("," ~ identifier ~ ":" ~ expr)* ~ "}"
              }
@@ -111,15 +111,17 @@ prefix = _{ not }
 
 call_begin = { "(" }
 
-//array_access = { identifier ~ "[" ~ expr ~ "]"}
-
 obj_initialize = { typ ~ "()" }
 call_args = { "(" ~ (expr ~ ("," ~ expr)*)? ~ ")" }
 func_call = { identifier ~ call_args }
 
 term = _{ literal | obj_initialize | func_call | identifier | "(" ~ expr ~ ")" | "[" ~ expr ~ "]" }
 
-expr = { prefix* ~ term ~ (binop ~ prefix* ~ term ) * }
+array_access = { "[" ~ expr ~ "]" }
+postfix = _{ array_access }
+term_ext = _ { prefix* ~ term ~ postfix* }
+
+expr = { term_ext ~ (binop ~ term_ext) * }
 
 
 /* ----- Statements ---- */

diff --git a/pollen/tests/degree.pollen b/pollen/tests/degree.pollen
@@ -0,0 +1,8 @@
+def degree() {
+  // Output is a (Segment, int) set.
+  parset out_degs[(Segment*int)];
+
+  for segment in in_g.segments {
+    emit (segment, segment.edges.size()) to out_degs;
+  }
+}
diff --git a/pollen/tests/flip.pollen b/pollen/tests/flip.pollen
@@ -0,0 +1,39 @@
+def flip() {
+  graph out_g;
+
+  for path in paths {
+    flip_path(path);
+  }
+}
+
+def flip_path(path: Path) {
+  parset out_steps[Step, out_g];
+  max_step_idx = path.length() - 1;
+  if is_rev(path) {
+    for step in path.steps {
+      emit
+      { step with handle: 
+          { step.handle with
+            orientation: !step.handle.orientation
+          },
+          idx: max_step_idx - step.idx
+      } to out_steps;
+    }
+  }
+}
+
+def is_rev(path: Path) {
+  fw = 0;
+  bw = 0;
+  for step in path.steps {
+    sh = step.handle;
+    len = sh.segment.length();
+    if sh.orientation {
+      fw = fw + len;
+    }
+    else {
+      bw = bw + len;
+    }
+  }
+  return bw > fw;
+}
diff --git a/pollen/tests/overlap.pollen b/pollen/tests/overlap.pollen
@@ -0,0 +1,14 @@
+def overlap() {
+  // Output is a (Path, Path) set.
+  parset out_overlaps[(Path*Path)];
+
+  for path in paths {
+    for step in path.steps {
+      for s in step.handle.segment.steps {
+        if s.path != path {
+          emit (path, s.path) to out_overlaps;
+        }
+      }
+    }
+  }
+}
diff --git a/pollen/tests/parse-out/crush.expect b/pollen/tests/parse-out/crush.expect
@@ -159,8 +159,10 @@ Prog {
                             },
                             EmitTo {
                                 expr: RecordUpdate {
-                                    parent: Id(
-                                        "segment",
+                                    parent: Var(
+                                        Id(
+                                            "segment",
+                                        ),
                                     ),
                                     fields: [
                                         RecordField {

diff --git a/pollen/tests/parse-out/degree.expect b/pollen/tests/parse-out/degree.expect
@@ -0,0 +1,76 @@
+Prog {
+    imports: [],
+    func_defs: [
+        FuncDef {
+            name: Id(
+                "degree",
+            ),
+            args: [],
+            ret_typ: None,
+            stmts: [
+                ParsetDecl {
+                    id: Id(
+                        "out_degs",
+                    ),
+                    typ: Tuple(
+                        Node,
+                        Int,
+                    ),
+                    graph_id: None,
+                },
+                For {
+                    id: Id(
+                        "segment",
+                    ),
+                    iterator: FieldAccess {
+                        object: Var(
+                            Id(
+                                "in_g",
+                            ),
+                        ),
+                        field: Var(
+                            Id(
+                                "segments",
+                            ),
+                        ),
+                    },
+                    body: Block {
+                        stmts: [
+                            EmitTo {
+                                expr: Tuple {
+                                    lhs: Var(
+                                        Id(
+                                            "segment",
+                                        ),
+                                    ),
+                                    rhs: MethodCall {
+                                        object: FieldAccess {
+                                            object: Var(
+                                                Id(
+                                                    "segment",
+                                                ),
+                                            ),
+                                            field: Var(
+                                                Id(
+                                                    "edges",
+                                                ),
+                                            ),
+                                        },
+                                        method: Id(
+                                            "size",
+                                        ),
+                                        args: [],
+                                    },
+                                },
+                                set_id: Id(
+                                    "out_degs",
+                                ),
+                            },
+                        ],
+                    },
+                },
+            ],
+            ret: None,
+        },
+    ],
+}