Add collection constructor operations and rewrite analyses to remove side effects

[jfeser]

Closes #715, #624

I did not pick up the changes from #680, although they're of a similar flavor.

Rewrites fvsof and sizesof in the jax module to remove their side-effects. This makes them compatible with memoization from #716.

fvsof's behavior changes for terms that contain collections: the collection construction operations now appear in the term free variables.

[eb8680]

fvsof's behavior changes for terms that contain collections: the collection construction operations now appear in the term free variables.

This seems like a big breaking change, are we sure it's necessary? I would expect lots of places here and downstream to depend on tests like if not fvsof(x): ... to gate partial evaluation behavior.

[jfeser]

No, it's not, but we should think about what semantics we want for fvsof. One way to view this change is that collections are a kind of shorthand for the real term, which is exactly an operation tree; in this view, fvsof should reflect that. Another view is that collection operations are an internal convenience that we want to hide from users, and should also be hidden from fvsof. I think the second view is the pragmatic one. Both are easy enough to implement.

[jfeser]

This PR now doesn't change the behavior of fvsof.

[eb8680]

How does this affect extensibility? With evaluate, defdata, nested_type etc the story is clear - a downstream user can register type-specific implementations that make them compatible with new third-party types. We don't want this new mechanism to break that and require edits to effectful to support new types.

I'd also worry that this is proliferating rather than consolidating traversal logic (#682 ) - isn't this morally duplicating evaluate's type-unpacking behavior? How do we avoid that, or at least enforce consistency across traversals?

[eb8680]

Should share a common base with CollectionConstrOperation

[eb8680]

This type rule is wrong/too weak, I would expect this to break typeof for tuples

[jfeser]

typeof runs CollectionConstrOperations rather than handling them. This allows us to keep the smarter type checking behavior from nested_type. I don't think we have rich enough types at the moment to type check the tuple constructor.

[eb8680]

I think this needs at least one type parameter for the return type

[eb8680]

I wonder if this should be a public classmethod of one of the new Operation classes. It seems like the right way to access these operations is to look them up by return type rather than by name, and to have an extension API that allows someone to register new ones as with to the singledispatch-based extensibility of evaluate and defdata.

[jfeser]

Sure.

[eb8680]

As I wrote in another comment, I think we probably want to look these up by type (in a singledispatch registry that is subtyping-aware) rather than by name to preserve our extensibility story

[jfeser]

These operations aren't the same as their corresponding constructors (e.g. tuple takes a single sequence argument, not a sequence of arguments). I think we only need a single sequence constructor though (probably as_tuple).

[jfeser]

evaluate extensions would need to use _as_type to allow uniform handling by analyses. This requires updating the extenders of evaluate, but I think it's better than the status quo.

Right now, typeof works because both evaluate and nested_type can be extended; nested_type unpacks the collection-typed analysis results produced by evaluate. An extension to evaluate generally requires an extension to nested_type. Some analyses, like fvsof, don't have the equivalent of nested_type, but can be implemented correctly using side effects. Other analyses, like Term.__str__, simply don't work correctly in the presence of extensions.

This PR doesn't change the status quo around nested_type, mainly because I think that's a substantial change, and I'm not sure we actually have rich enough type annotations to capture what nested_type does.

Re, traversal logic, I think this is a simplification because more analyses can become pure handlers of apply (as is done in this PR for fvsof and sizesof). I'm not totally sure what it would be mean for interpretations to be consistent about their traversal behavior beyond consuming the term tree in the same order.