feat(ir): add list comprehension to IR #136

tohrnii · 2023-02-07T13:28:57Z

Partly addressing #87. This PR adds list comprehension and list folding support to the IR.

bobbinth

Thank you! Looks good! This is not a full review yet - but I left some comments inline.

ir/src/constraints/list_comprehension.rs

bobbinth · 2023-02-15T08:06:15Z

ir/src/constraints/list_comprehension.rs

+/// Returns an error if there is an error while parsing the sub-expression.
+fn parse_lc_expr(
+    expression: &Expression,
+    lc: &ListComprehension,


It seems to me that we are passing ListComprehension around to get access to its context and then use loops to figure out which member corresponds to which iterable. If so, a better solution could be to build a map from identifier to iterable upfront (maybe put this map into something a struct - something like IterableContext), and then pass this map around instead ListComprehension.

This way, instead of doing something like:

for (member, iterable) in lc.context() { if ident == member { match iterable {

You'd be able to do something like:

let iterable = context.get_iterable(ident)?; match iterable {

You may be able to take it even a step further and put IdentifierType for each member into the context. This way, you won't need to pass symbol_table around either.

Thanks, done. I moved it to just a BTreeMap<Identifier, Iterable>. I didn't remove symbol table as it is required to build context for list folding inside list comprehension.

bobbinth · 2023-02-15T08:13:34Z

ir/src/constraints/list_comprehension.rs

+                            return Ok(Expression::NamedTraceAccess(NamedTraceAccess::new(
+                                ident.clone(),
+                                range.start() + i,
+                                CURRENT_ROW,
+                            )));


Why do we hard-code CURRENT_ROW here? Shouldn't the offset come from trace_columns?

The AST has current row trace columns stored as general element since there's no way to disambiguate between whether it's a trace column or any other symbol type like an intermediate variable. So, if it's stored as a trace column in symbol table, we know that it should have 0 row offset here.

I might be missing something here, but how is something like this stored in the symbol table:

trace_columns: main: [a, b] integrity_constraints: let x = a' enf x = b + 2

The trace columns are stored as:

TraceColumns(TraceColumns)

where

pub struct TraceColumns { trace_segment: TraceSegment, offset: usize, size: usize, }

where offset is the column offset of the trace column group instead of row offset.

When we insert the intermediate variable let x = a', a' would be stored in the graph as:

TraceElement(IndexedTraceAccess { trace_segment: 0, col_idx: 0, row_offset: 1 })

I think the confusion might be that a and a' are stored differently in the AST. When we encounter a at the parser level we don't know if it's a trace column or any other identifier so we store it as Expression::Elem, however when we encounter a' at the parser level, we know that it should be referring it to a trace column so we store it as Expression::NamedTraceAccess(NamedTraceAccess) where

pub struct NamedTraceAccess { name: Identifier, idx: usize, row_offset: usize, }

Maybe there's a better way to do this since this is certainly a little confusing. Please let me know if you have other suggestion.

So, when we parse the AST expression while unfolding list comprehension, if the expression is of type Elem, we know that it should be referring to current row.

ir/src/constraints/list_comprehension.rs

bobbinth · 2023-02-15T08:15:09Z

ir/src/constraints/list_comprehension.rs

+                            return Ok(Expression::NamedTraceAccess(NamedTraceAccess::new(
+                                ident.clone(),
+                                i,
+                                NEXT_ROW,
+                            )));


Why do we hard-code NEXT_ROW here? Shouldn't the offset come from trace_columns?

You are right, it should come from named_trace_access here. Thanks.

bobbinth · 2023-02-15T08:15:18Z

ir/src/constraints/list_comprehension.rs

+                            return Ok(Expression::NamedTraceAccess(NamedTraceAccess::new(
+                                ident.clone(),
+                                range.start() + i,
+                                NEXT_ROW,
+                            )));


Why do we hard-code NEXT_ROW here? Shouldn't the offset come from trace_columns?

You are right, it should come from named_trace_access here. Thanks.

ir/src/constraints/graph.rs

bobbinth

Thank you! Still not a full review - but I left some comments inline.

I think the biggest issue is that the constraint test file (list_comprehension.air) seems to be compiled incorrectly into Rust code. I'm not sure if this is due to this PR or is a prior issue - but I think it would be difficult to check the behavior until this is corrected.

ir/src/constraints/list_comprehension.rs

bobbinth · 2023-02-15T21:15:53Z

ir/src/constraints/list_comprehension.rs

+    let iterable = iterable_context.get(ident);
+    if let Some(iterable_type) = iterable {
+        match iterable_type {


If iterable_context.get() does not return Some, should this not be an error? As I mentioned in a comment from the other review, it might be good to create a dedicate struct for IterableContext and then we can do something like:

let iterable_type = iterable_context.get(ident)?; match iterable_type {

@bobbinth I think we can have symbols in a list comprehension expression that are not in the iterables of the list comprehension. For example, the following should be valid:

trace_columns: main: [c[3]] let x = 1 let a = [b + x for b in c]

Here, when we are unrolling the list comprehension, we just leave x alone and return it if we don't find a corresponding iterable. So the above list comprehension unrolls into:

[c[0] + x, c[1] + x, c[2] + x]

I'm assuming the only reason you wanted to go with a struct here was to keep returning error clean? If so, should we still move this to a dedicated struct?

Ah - makes sense! I'd probably add some comments and maybe restructure the function to make it more clear.

For example, for parse_elem() do we really need to pass in expression? If I understood it correctly, passing both ident and expression is redundant. And the only reason we pass it in is to clone it at the end. But we could have just as easily done something like Expression::Elem(ident) right?

And the logic of the whole function could be:

If the ident is not in the iterable_context, then just return Expression::Elem(ident) (and I guess any invalid references will be caught later).

Otherwise, process the ident as we do now.

bobbinth · 2023-02-15T21:19:02Z

ir/src/constraints/mod.rs

@@ -232,7 +236,15 @@ impl Constraints {
                self.insert_constraint(symbol_table, lhs, rhs)?
            }
            IntegrityStmt::Variable(variable) => {
-                symbol_table.insert_integrity_variable(variable)?
+                if let VariableType::ListComprehension(list_comprehension) = variable.value() {
+                    let vector = unfold_lc(list_comprehension, symbol_table)?;


nit: would a more appropriate name here be expressions rather than vector?

expressions seems a little vague to me since we are returning a vector of expressions to be used for a vector type variable in this context. Do you think variable_value is more accurate? Not a strong opinion though, I'm happy to go with expressions if that seems better to you.

I probably prefer expressions as well, but you could use expr_vector to capture both ideas

bobbinth · 2023-02-15T21:39:30Z

air-script/tests/list_comprehension/list_comprehension.rs

+    fn evaluate_aux_transition<F, E>(&self, main_frame: &EvaluationFrame<F>, aux_frame: &EvaluationFrame<E>, _periodic_values: &[F], aux_rand_elements: &AuxTraceRandElements<E>, result: &mut [E])
+    where F: FieldElement<BaseField = Felt>,
+          E: FieldElement<BaseField = Felt> + ExtensionOf<F>,
+    {
+        let current = aux_frame.current();
+        let next = aux_frame.next();
+        result[0] = current[0] - (((E::from(2_u64)).exp(E::PositiveInteger::from(3_u64))) * (current[7]));
+        result[1] = current[0] - (next[4] - (next[8]));
+        result[2] = current[0] - (current[5] - (current[10]));
+        result[3] = current[0] - (E::from(0_u64) + current[1] - (current[4]) - (current[8]) + E::from(1_u64) + current[2] - (current[5]) - (current[9]) + E::from(2_u64) + current[3] - (current[6]) - (current[10]));
+        result[4] = current[0] - ((current[4]) * (current[8]) + (current[5]) * (current[9]) + (current[6]) * (current[10]) + (current[7]) * (current[11]));
+        result[5] = current[0] - (current[1] + (current[4]) * (current[8]) + (current[5]) * (current[9]) + (current[6]) * (current[10]) + (current[7]) * (current[11]));
+    }


There is something wrong here and I'm not sure if it has to do with list comprehension or with something else. The main issue is that most constraints in the test file use clk column which comes from the main trace, but the constraints here use are all placed against current[0] which refers to column a in the auxiliary trace.

Thanks for catching this. This seems like it was an existing issue. I'll open a PR today with a fix.

I've opened #148.

bobbinth · 2023-02-15T21:45:00Z

air-script/tests/list_comprehension/list_comprehension.rs

+    {
+        let current = aux_frame.current();
+        let next = aux_frame.next();
+        result[0] = current[0] - (((E::from(2_u64)).exp(E::PositiveInteger::from(3_u64))) * (current[7]));


This is probably for the future but the exponentiation here should have been done as E::from(2_u64.pow(3)) instead of E::from(2_u64)).exp(E::PositiveInteger::from(3_u64). Basically, raising a constant to a constant power should result in a constant.

Actually, I would even say that the expression should have been just E::from(8_u64) but maybe this is an optimization step.

Thanks. Yeah, I think the second one should probably be an optimization step. I've added #149 to track these.

bobbinth · 2023-02-15T21:53:16Z

air-script/tests/list_comprehension/list_comprehension.air

+    let enumerate = [2^i * c for (i, c) in (0..4, c)]
+    enf clk = enumerate[3]
+
+    let diff0 = [x' - y' for (x, y) in (c, d)]
+    enf clk = diff0[0]
+
+    let diff1 = [x - y for (x, y) in (c[0..2], d[1..3])]
+    enf clk = diff1[1]
+
+    let x0 = [w + x - y - z for (w, x, y, z) in (0..3, b, c[0..3], d[0..3])]
+    enf a = x0[0] + x0[1] + x0[2]
+
+    let x1 = sum([c * d for (c, d) in (c, d)])
+    let y1 = prod([c + d for (c, d) in (c, d)])    
+    enf clk = x1
+
+    let x2 = sum([c * d for (c, d) in (c, d)])
+    let y2 = [m + x2 for m in fmp]
+    enf clk = y2[0]


This may be for a different PR, but most constraints here should not be allowed. For example, the first constraint actually does something like this:

enf clk = 2^3 * c[3]

But clk is from main trace (and thus has type of E::BaseField) while c[3] is from the auxiliary trace and has type of E. In practice this constraint will never be satisfied and ideally we should try to flag this as error.

Makes sense. I was a little sloppy in writing these examples (just wanted to write down different kinds of expressions possible. I'll improve on them to only include valid expressions.

But clk is from main trace (and thus has type of E::BaseField) while c[3] is from the auxiliary trace and has type of E. In practice this constraint will never be satisfied and ideally we should try to flag this as error.

We already have an issue #47 to track this.

bobbinth · 2023-02-15T22:04:06Z

ir/src/constraints/list_comprehension.rs

+                            return Ok(Expression::NamedTraceAccess(NamedTraceAccess::new(
+                                ident.clone(),
+                                range.start() + i,
+                                CURRENT_ROW,
+                            )));


I might be missing something here, but how is something like this stored in the symbol table:

trace_columns: main: [a, b] integrity_constraints: let x = a' enf x = b + 2

grjte

This is a partial review. I need to spend more time looking at the list_comprehension module.

As a general comment, I think it would be easier to read this file if the doc comments were more detailed and/or provided examples - in most cases, the doc comments explain what's happening in terms of internal terminology that we've chosen or refer to undocumented structs, so it's a bit hard to understand one thing without already understanding everything. Ideally, everything that's happening could be understood by just reading the function doc comments, and I think that's not really the case right now (or not easily)

ir/src/constraints/mod.rs

ir/src/constraints/list_comprehension.rs

ir/src/constraints/graph.rs

ir/src/constraints/list_comprehension.rs

tohrnii · 2023-02-23T04:05:20Z

@grjte I've addressed your comments except adding more detailed docs. I'm working on them now.

One question is if we should rename unfold_lc to something like expand_lc (or something else) because unfold might be confusing since we also have list folding?

grjte

Thanks @tohrnii I think it's looking really good. The comments you've added are very helpful. It feels quite clear to read through it at this point. I left several things inline, but they're all minor.

One final note is that I feel like we might be able to improve the pattern for the list_comprehension module using a struct + methods instead of all of the independent functions. I think it would be hard to do this now, but that it might be worth a look after #150 (or maybe as part of it)

ir/src/constraints/constraint.rs

ir/src/constraints/list_comprehension.rs

ir/src/tests/list_comprehension.rs

tohrnii · 2023-02-24T22:37:19Z

Thank you @bobbinth and @grjte for the reviews. Made suggested changes.

grjte

LGTM, great work @tohrnii, thanks! I left 2 final non-blocking nits inline

grjte · 2023-02-25T12:18:45Z

air-script/tests/list_comprehension/list_comprehension.rs

+use winter_math::fields::f64::BaseElement as Felt;
+use winter_math::{ExtensionOf, FieldElement};
+use winter_utils::collections::Vec;
+use winter_utils::{ByteWriter, Serializable};


nit: these could be condensed into 2 lines

I'm not sure how to do that. The import function takes a module path and type. So, currently we do:

scope.import("winter_utils::collections", "Vec")

If I move collections to the type argument:

scope.import("winter_utils", "collections::Vec")

this generates output:

use winter_utils::{collections, ByteWriter, Serializable};

This is what I had done earlier to combine the imports before @bobbinth pointed out the problem.

Looking at the import constructor it just joins the two strings:

pub fn new(path: &str, ty: &str) -> Self { Import { line: format!("{}::{}", path, ty), vis: None, } }

So, will have to look into it more. If you have any ideas on how to achieve this, please let me know otherwise I'll dig into this more later.

Ah yeah, makes sense. It's not worth any time before the next version, so let's just leave it.

ir/src/constraints/list_comprehension.rs

bobbinth

All looks good! Thank you!

tohrnii force-pushed the tohrnii-lc-ir branch from 3971393 to dbc4f6e Compare February 7, 2023 14:42

tohrnii mentioned this pull request Feb 7, 2023

Add list comprehension to IR #133

Closed

tohrnii force-pushed the tohrnii-lc-ir branch 7 times, most recently from 6e02891 to 54b5b88 Compare February 8, 2023 03:32

tohrnii marked this pull request as ready for review February 8, 2023 03:37

tohrnii requested review from grjte, Fumuran and Al-Kindi-0 February 8, 2023 03:37

tohrnii force-pushed the tohrnii-lc-ir branch 5 times, most recently from f45ea4a to 07906a1 Compare February 8, 2023 15:01

tohrnii mentioned this pull request Feb 11, 2023

Tracking PR for v0.2 release #79

Merged

14 tasks

tohrnii requested review from bobbinth and removed request for grjte February 13, 2023 15:11

tohrnii force-pushed the tohrnii-lc-ir branch from 07906a1 to 96cc68a Compare February 13, 2023 15:38

bobbinth requested changes Feb 15, 2023

View reviewed changes

tohrnii force-pushed the tohrnii-lc-ir branch 4 times, most recently from 51529a6 to 9635928 Compare February 15, 2023 18:33

tohrnii requested a review from bobbinth February 15, 2023 18:35

bobbinth requested changes Feb 15, 2023

View reviewed changes

tohrnii force-pushed the tohrnii-lc-ir branch from 9635928 to 4d01e75 Compare February 16, 2023 12:45

tohrnii force-pushed the tohrnii-lc-ir branch 4 times, most recently from 1f68e8e to 925c89c Compare February 17, 2023 11:37

tohrnii requested a review from bobbinth February 17, 2023 11:37

tohrnii force-pushed the tohrnii-lc-ir branch from 925c89c to 818a57e Compare February 17, 2023 11:41

grjte self-requested a review February 20, 2023 10:38

grjte suggested changes Feb 20, 2023

View reviewed changes

tohrnii force-pushed the tohrnii-lc-ir branch from 0fd08c3 to 638d8e1 Compare February 23, 2023 04:02

tohrnii requested a review from grjte February 23, 2023 04:05

tohrnii force-pushed the tohrnii-lc-ir branch 2 times, most recently from 471f170 to 7fee3bd Compare February 24, 2023 04:04

grjte suggested changes Feb 24, 2023

View reviewed changes

tohrnii added 4 commits February 24, 2023 22:42

feat(ir): add list comprehension to IR

03e8638

feat(ir): add support for list folding expressions

fcddeea

test: add list comprehension IR tests

4b71bff

test: add list comprehension integration test

92cbdd2

tohrnii force-pushed the tohrnii-lc-ir branch from f6752b2 to a758145 Compare February 24, 2023 22:43

tohrnii requested a review from grjte February 24, 2023 22:43

tohrnii mentioned this pull request Feb 25, 2023

feat(docs): add docs for list comprehension #173

Merged

grjte approved these changes Feb 25, 2023

View reviewed changes

refactor(ir): improve error messages and other minor refactors

e3ef487

tohrnii force-pushed the tohrnii-lc-ir branch from a758145 to e3ef487 Compare February 25, 2023 13:03

bobbinth approved these changes Feb 25, 2023

View reviewed changes

grjte merged commit 9a0b438 into next Feb 27, 2023

grjte deleted the tohrnii-lc-ir branch February 27, 2023 09:43

grjte mentioned this pull request Feb 27, 2023

Tracking PR for v0.3 release #175

Merged

tohrnii mentioned this pull request Mar 7, 2023

List Comprehension #87

Closed

feat(ir): add list comprehension to IR #136

feat(ir): add list comprehension to IR #136

Conversation

tohrnii commented Feb 7, 2023 • edited Loading

bobbinth left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tohrnii Feb 15, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tohrnii Feb 16, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bobbinth left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tohrnii Feb 15, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tohrnii Feb 16, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tohrnii Feb 16, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

grjte left a comment • edited Loading

Choose a reason for hiding this comment

tohrnii commented Feb 23, 2023

grjte left a comment

Choose a reason for hiding this comment

tohrnii commented Feb 24, 2023

grjte left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tohrnii Feb 25, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bobbinth left a comment

Choose a reason for hiding this comment

tohrnii commented Feb 7, 2023 •

edited

Loading

tohrnii Feb 15, 2023 •

edited

Loading

tohrnii Feb 16, 2023 •

edited

Loading

tohrnii Feb 15, 2023 •

edited

Loading

tohrnii Feb 16, 2023 •

edited

Loading

tohrnii Feb 16, 2023 •

edited

Loading

grjte left a comment •

edited

Loading

tohrnii Feb 25, 2023 •

edited

Loading