Rust constr enum #1941

Schaeff · 2024-10-23T14:07:59Z

Identity is currently a struct which acts as an union: all kinds of identities are represented using the same struct fields, with runtime checks that the right fields are used.

In the context of #1934 where we add new kinds of identities, it seems advantageous to move to an enum on the Rust side as well.

Todo:

reimplement ids
minimize code duplication between lookups and permutations, if possible
remove set_id
find better way to compare identities
fix connect

Schaeff · 2024-10-23T15:36:48Z

ast/src/analyzed/mod.rs

+    }
+}
+
+// TODO This is the only version of Identity left.


@chriseth can I remove this comment?

The idea behind this comment was to refactor Identity so that it is only generic over the field, and not about what a selected expressions or an expression / algebraic expression is. If you did this change, then you can remove it.

Oh I didnt get that but yes that's exactly what I did :)

Can you remove the TODO?

Schaeff · 2024-10-23T16:43:39Z

executor/src/witgen/machines/machine_extractor.rs

@@ -309,6 +336,9 @@ fn all_row_connected_witnesses<T>(
                        witnesses.extend(in_rhs);
                    }
                }
+                Identity::Permutation(..) | Identity::Connect(..) => {
+                    todo!("same as plookup, avoid repetition?")
+                }


The comment of this function does not mention connect, is this something we use?

Connect is very sporadically implemented.

Schaeff · 2024-10-23T17:14:43Z

executor/src/witgen/mod.rs

@@ -227,7 +228,13 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
            // as they are assumed to be handled in stage 0.
            let polynomial_identities = identities


Is it possible for identities to still contains anything else than polynomial identities? If so, due to the filtering these would get ignored? Or do we not have any non-polynomial identities, in which case this can be a hard cast rather than a filter?

I think the filtering here could be improved in general. I don't see polynomial constraints between lower-stage witness column being filtered out, for example. I also don't see a reason why this should only contain polynomial constraints at this stage, e.g. we shouldn't filter out the lookups. Of course, usually when we use a multi-stage approach the idea is to remove the lookups, but conceptually, we could still have a lookup at a higher stage (although the backend probably would not support it).

How about we leave this as is in this PR?

chriseth · 2024-10-24T08:03:48Z

ast/src/analyzed/mod.rs

+        }
+    }
+
+    pub fn source(&self) -> &SourceRef {


There is a trait called SourceReference, I think we should implement that instead.

chriseth · 2024-10-24T08:09:45Z

executor/src/witgen/machines/mod.rs

+    Permutation,
+}
+
+impl<'a, T> ConnectingIdentityRef<'a, T> {


Is there any way we could re-write the code that uses these Refs so that they are not needed?

Hmm.. So assuming we only rely on the overlap of the two (left and right selected expressions) we could just pass that, but I don't really find that cleaner and we lose the ability to do things specifically for each case.

Also this assumption doesn't seem to hold since we have code like this for WriteOnceMemory

if !parts.connecting_identities.values().all(|i| i.is_lookup()) { return None; }

Or this for BlockMachine

// Connecting identities should either all be permutations or all lookups. let connection_type = connecting_identities .values() .map(|id| id.kind()) .unique() .exactly_one() .ok()?;

If your point is that we could refactor these, maybe? I think the current way turns what was a runtime invariant into a compile-time one, with minimal changes to the rest of the code.

I think we are not talking about the same thing. My question was if we can avoid adding a ..Ref type for every type.

Ok then I don't understand, there's a single Ref type and it's ConnectingIdentityRef. It exists because connecting identities can only be of certain types, so using &'a Identity would lead to a bunch of unreachable!().

Ah sorry, I thought this was the code used in path canonicalizer which is also sprinkled with ...Refs.

I don't really like it either to be honest.

What about

struct ConnectingIdentity<'a, T> { left: &'a SelectedExpressions<T>, right: &'a SelectedExpressions<T>, kind: ConnectingIdentityKind }

Where kind is Lookup or Permutation and will not need to be extended for PhantomLookup and PhantomPermutation?

Then we can remove some duplicated code and still rely on kind for the cases I pointed out above?
I don't really see how traits help here, since at the end of the day we need to be able to store either of these two possibilities anywhere, so we'd have to use dyn ConnectingIdentity which seems unnatural here due to the fixed number of variants. But maybe I'm not zooming out enough here.

I tried the above and I like it more

Cool! Don't we need the ID? Oh, it's in a hash map from id anyway, right?

Yeah I put the id for now but it did seem redundant, let's see.

chriseth · 2024-10-24T08:15:08Z

pilopt/src/lib.rs

+    struct CanonicalIdentity<'a, T>(&'a Identity<T>);
+
+    impl<T: PartialEq> PartialEq for CanonicalIdentity<'_, T> {
+        fn eq(&self, other: &Self) -> bool {


At least this one can be implemented using PartialOrd, right?

Excellent point

Done using Ord

Schaeff · 2024-10-24T08:58:50Z

Reverting to draft as I realised I broke connect

Schaeff · 2024-10-24T09:13:06Z

pilopt/src/lib.rs

    // Set of (left, right) tuples.
    let mut identity_expressions = BTreeSet::new();
    let to_remove = pil_file
        .identities
        .iter()
        .enumerate()
        .filter_map(|(index, identity)| {
-            match identity_expressions.insert((&identity.left, &identity.right)) {


By the way, this seems like a bug, as these would be considered duplicates:

[a] in [b] [a] connect [b] [a] is [b]

I don't see that typically happening but still.

chriseth · 2024-10-24T16:25:36Z

pilopt/src/lib.rs

+                    ) => a.cmp(b),
+                    (
+                        Identity::Lookup(LookupIdentity {
+                            left: a, right: b, ..


You could use a generic left and right function for any non-polynomial identity.

I thought about it but given that these are in different variants, I don't see how we could use it:
now we have:

Identity::Lookup(LookupIdentity { left, right, .. }) | Identity::Permutation(PermutationIdentity { left, right, ..}) => { left }

With a trait, we can't treat both variants at the same time because we can only bind to variables of the same type in patterns:

// does NOT work Identity::Lookup(i) | Identity::Permutation(i) => { i.left()

For this to be possible, we need to use the same struct inside the variant:

Identity::Lookup(ConnectingIdentity<Lookup>), Identity::Permutation(ConnectingIdentity<Permutation>),

... but this doesn't work either because the generic types also need to match!

So we go one step further and use the same variant

Identity::Connection(ConnectingIdentity)

And store the kind inside the ConnectingIdentity.

This is nice but my issue is that at this point we drifted away completely from the PIL enum.

My conclusion is that relying on patterns to work across connecting identities, while verbose, is fine as it preserves the structure of the PIL enum. If we refactor the PIL enum to be nested, then the nested approach would work in Rust as well. I don't think this refactor is necessarily beneficial (because of its nested nature), nor a priority.

chriseth · 2024-10-24T16:27:31Z

executor/src/witgen/machines/mod.rs

+    Permutation(&'a PermutationIdentity<T>),
+}
+
+impl<'a, T: Display> Display for ConnectingIdentityRef<'a, T> {


Could use derive_more here.

chriseth · 2024-10-24T16:33:20Z

executor/src/witgen/machines/machine_extractor.rs

                // For lookups, any lookup calling from the current machine belongs
                // to the machine; lookups to the machine do not.
-                let all_refs = &refs_in_selected_expressions(&i.left) & (&all_witnesses);
+                let all_refs = match i {
+                    Identity::Polynomial(identity) => {


Maybe extract that into a function?

done, with all_children

That was wrong so done, just extracting as suggested

chriseth · 2024-10-24T16:33:34Z

executor/src/witgen/machines/machine_extractor.rs

+            .filter_map(|i| {
+                let id = i.id();
+                match i {
+                    Identity::Polynomial(_) => None,


Maybe extract that into a function?

simplified due to refactor of ConnectingIdentity

Schaeff · 2024-10-24T23:15:43Z

This is ready for another review @chriseth

chriseth · 2024-10-25T10:28:25Z

ast/src/analyzed/mod.rs

+}
+
+impl<T> PolynomialIdentity<T> {
+    pub fn new(id: u64, source: SourceRef, expression: AlgebraicExpression<T>) -> Self {


Do we acutally need this when the fields are all pub?

No, personal preference but happy to remove these especially since we rely heavily on destructuring anyway

chriseth · 2024-10-25T10:39:59Z

executor/src/witgen/machines/mod.rs

@@ -151,6 +150,67 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
    }
 }

+#[derive(Clone, Copy)]
+pub struct ConnectingIdentity<'a, T> {


Maybe one line of documentation?

Hm, not sure I would agree to this. A polynomial identity can also cross machine boundaries (depending on the definition, it would be the same machine, but still...). Shouldn't we be more explicit about the kind? Or the fact that it has two parts that do not need to match on the same row or something like that if you want to be abstract?

I would just be explicit and say, lookup or permutation.

chriseth

I still think the doc is confusing (people will wonder if it is related to Identity or not), but ok.

Schaeff added 4 commits October 22, 2024 15:39

wip

778036c

wip

071b357

refactor, make compile

4b0486a

clean, remove unnecessary generics

d300846

Schaeff commented Oct 23, 2024

View reviewed changes

self reviewed all except witgen

ac7000c

Schaeff commented Oct 23, 2024

View reviewed changes

self review witgen

e5bde54

Schaeff marked this pull request as ready for review October 23, 2024 17:23

Schaeff added 2 commits October 23, 2024 19:53

fix tests

2b52b8d

lint

3afbafe

chriseth reviewed Oct 24, 2024

View reviewed changes

Schaeff added 2 commits October 24, 2024 10:18

use SourceReference trait

0989a1e

define all traits on CanonicalIdentity based on Ord

d0d602b

Schaeff marked this pull request as draft October 24, 2024 08:58

fmt

7d862fb

Schaeff commented Oct 24, 2024

View reviewed changes

Schaeff added 2 commits October 24, 2024 16:16

fix connect

8bc60ef

remove set_id

79946e0

Schaeff marked this pull request as ready for review October 24, 2024 15:45

Merge branch 'main' of github.com:powdr-labs/powdr into rust-constr-enum

7803534

chriseth reviewed Oct 24, 2024

View reviewed changes

address review comments

bf7728e

Schaeff added 2 commits October 25, 2024 00:02

fix extraction

b60d71f

rename type parameter

bf2a796

chriseth reviewed Oct 25, 2024

View reviewed changes

Schaeff added 3 commits October 25, 2024 13:21

address review comments

e7abde2

rename to Connection, fix docs

7382fee

rename

208971b

chriseth approved these changes Oct 25, 2024

View reviewed changes

Schaeff added this pull request to the merge queue Oct 25, 2024

Merged via the queue into main with commit 3617c90 Oct 25, 2024
14 checks passed

Schaeff deleted the rust-constr-enum branch October 25, 2024 15:00

		@@ -227,7 +228,13 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
		// as they are assumed to be handled in stage 0.
		let polynomial_identities = identities

Rust constr enum #1941

Rust constr enum #1941

Conversation

Schaeff commented Oct 23, 2024 • 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

Schaeff Oct 23, 2024 • 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

Schaeff Oct 24, 2024 • 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

Schaeff Oct 24, 2024 • 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

chriseth Oct 24, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Schaeff commented Oct 24, 2024

Schaeff Oct 24, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Schaeff Oct 24, 2024 • 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

Schaeff commented Oct 24, 2024

Choose a reason for hiding this comment

Schaeff Oct 25, 2024 • 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

chriseth left a comment

Choose a reason for hiding this comment

Schaeff commented Oct 23, 2024 •

edited

Loading

Schaeff Oct 23, 2024 •

edited

Loading

Schaeff Oct 24, 2024 •

edited

Loading

Schaeff Oct 24, 2024 •

edited

Loading

chriseth Oct 24, 2024 •

edited

Loading

Schaeff Oct 24, 2024 •

edited

Loading

Schaeff Oct 24, 2024 •

edited

Loading

Schaeff Oct 25, 2024 •

edited

Loading