Adding type annotations to intermediate representation #320

yihozhang · 2023-12-08T05:19:59Z

This PR implements some important refactoring. At a high level, this PR gets rid of many legacy code and use a set of generic IRs which allows greater code sharing (thus reduced codebase size).

Fully remove the old NormExpr and NormAction, use the new CoreIR instead.
Generalizes front-end IR (e.g., Rule) and backend-IR (e.g., CoreRule) to be parametric over the head and leaf symbols.
Add type resolution so that typechecking returns a type-annotated expression.
Replace the type checker for actions to use the constraint-based one.
Refactor the backend Action->Program compiler so that it no longer does unnecessary type checking and works with CoreAction instead of Action.
desugar and resugar flags are replaced with a --show desugared-egglog flag, which prints an egglog program produced by the desugaring phase.

The current compilation workflow is

Command -[desugar]-> NCommand: does some lightweight desugaring, e.g., turns a function into a relation
Rule in NCommand -[lower]-> CoreRule: lowering to CoreRule for typechecking
(Rule, CoreRule) -[typecheck]-> ResolvedRule: ResolvedRule is a Rule with all its function symbols and variables type annotated.
Other planned transformations (e.g. proofs) over ResolvedRule
ResolvedRule -[lower]-> ResolvedCoreRule
ResolvedCoreRule -[compilation]-> Program

Fixes #80, fixes #113, fixes #196.

However, the type constraint part is almost done

…nction generically.

2. expr_to_norm_actions -> expr.to_norm_actions

Some remaining work in main.rs, test.rs, to_resolved, and others.

* Change the specification of accept * Use indexmaps more to enforce a consistent variable ordering for Load::Subst * Refactor `compile_exprs`

Fix clippy errors and extraction in yihong's typechecking branch

…gglog into yihozhang-type-annotation

…glog * Remove should_run flag

…gglog into yihozhang-type-annotation

src/ast/desugar.rs

ezrosent

Minor comments (this is the first round. I think there's more to do but this may be a good place to start)

ezrosent · 2024-01-12T20:22:55Z

src/ast/mod.rs

-
-#[derive(Debug, Clone, Eq, PartialEq, Hash)]
-pub enum NCommand {
+pub enum GenericNCommand<Head, Leaf, Ann> {


While the context is fresh, let's add a doc comment explaining what a GenericNCommand is.

Also, offline discussed that different names would be better here.

optional: Let's think about naming things Generic.

ezrosent · 2024-01-12T20:24:33Z

src/ast/mod.rs

@@ -40,232 +43,110 @@ impl Display for Id {
    }
 }

-pub type CommandId = usize;
+pub type NCommand = GenericNCommand<Symbol, Symbol, ()>;


More descriptive name and/or doc

(SourceCommand? UntypedCommand? Happy to have the one that the user sees just be Command or similar, but then we'll want good docs explaining the relationship that other commands (e.g. ResolvedCommand have to Command).

Done. Same above: I will make a separate PR to use the DCommand/RawDCommand/ResolvedDCommand scheme.

ezrosent · 2024-01-12T20:28:45Z

src/core.rs

@@ -0,0 +1,755 @@
+/// This file implements the core IR of the language, which is called CoreRule.


Doc comments for a module should have //! as their prefix.

ezrosent · 2024-01-12T20:32:27Z

src/core.rs

+        ))
+    }
+}
+


optional: consider shorter names for the return types in to_query:

I added a MappedExpr type (with brief documentation) as a shorthand

ezrosent · 2024-01-12T20:34:24Z

src/core.rs

+}
+
+#[derive(Debug, Clone)]
+pub struct GenericCoreRule<BodyF, HeadF, Leaf> {


We could use a doc comment here (perhaps something that sheds light on the F at the end).

ezrosent · 2024-01-15T07:39:24Z

src/ast/mod.rs

-                }
-                var
-            });
+    pub fn map_def_use(&self, fvar: &mut impl FnMut(&Leaf, bool) -> Leaf) -> Self {


Small nits:

Why is this pub and not pub(crate) ?

Document this method.

It is actually funny that every method in this impl are dead code if not pub. I left a note here. I think map_def_use is referred to in term.rs so it might be needed in the near future. All these functions are either from the very original codebase or from the NormFact refactor. I'm also happy to just delete all of them.

If we can delete them then I think we should :)

Up to you though

ezrosent · 2024-01-15T07:45:24Z

src/constraint.rs

+        rule: &CoreRule,
+        typeinfo: &TypeInfo,
+    ) -> Result<(), TypeError> {
+        let CoreRule { head, body } = rule;


optional: you can move this to the function parameter, I think:

CoreRule { head, body } : &CoreRule

ezrosent · 2024-01-15T07:48:22Z

src/core.rs

+        }
+    }
+
+    // Different from `from_resolution`, this function only considers function types and not primitives.


I think this comment can be a /// comment

Why is it pub ?

Done; I changed them to pub(crate)

ezrosent · 2024-01-15T07:49:15Z

src/core.rs

+        types: &[ArcSort],
+        typeinfo: &TypeInfo,
+    ) -> Option<ResolvedCall> {
+        if let Some(ty) = typeinfo.func_types.get(head) {


nit, this can be:

let ty = typeinfo.func_types.get(head)?; let expected = ...

ezrosent · 2024-01-15T07:49:49Z

src/core.rs

+            }
+        }
+
+        assert!(resolved_call.len() == 1);


assert_eq!(resolved_call.len(), 1)

yihozhang

@ezrosent Thanks for the awesome review. The review is super helpful for improving this PR and I really appreciate it. I should have addressed most of your comments.

yihozhang · 2024-01-18T23:28:12Z

src/actions.rs

+    Union(usize),
+    Extract(usize),
+    Panic(String),
+}


yihozhang · 2024-01-18T23:29:26Z

src/actions.rs

+            instructions: Vec::new(),
+        };
+
+        for a in actions.0.iter() {


You are absolutely right (also changed a few other places)

yihozhang · 2024-01-18T23:30:33Z

src/ast/desugar.rs

    // Store the parser because it takes some time
    // on startup for some reason
    parser: ast::parse::ProgramParser,
-    pub(crate) expr_parser: ast::parse::ExprParser,
-    pub(crate) action_parser: ast::parse::ActionParser,
+    // yz (dec 5): Comment out since they are only used in terms.rs which are deleted


I think leaving them here makes Oliver's life easier later when he wants to add terms back, but I'm open to deleting them.

yihozhang · 2024-01-19T02:37:35Z

src/ast/expr.rs

+    }
+}
+
+pub type Expr = GenericExpr<Symbol, Symbol, ()>;


yihozhang · 2024-01-19T02:38:53Z

src/ast/expr.rs

+pub(crate) type ResolvedExpr = GenericExpr<ResolvedCall, ResolvedVar, ()>;
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
+pub enum GenericExpr<Head, Leaf, Ann> {


Expr/RawExpr/ResolvedExpr seems like a good idea. I will make a separate PR for the renaming.

yihozhang · 2024-01-19T03:19:56Z

src/constraint.rs

+        rule: &CoreRule,
+        typeinfo: &TypeInfo,
+    ) -> Result<(), TypeError> {
+        let CoreRule { head, body } = rule;


yihozhang · 2024-01-19T03:22:38Z

src/core.rs

+        }
+    }
+
+    // Different from `from_resolution`, this function only considers function types and not primitives.


Done; I changed them to pub(crate)

yihozhang · 2024-01-19T03:23:26Z

src/core.rs

+        types: &[ArcSort],
+        typeinfo: &TypeInfo,
+    ) -> Option<ResolvedCall> {
+        if let Some(ty) = typeinfo.func_types.get(head) {


yihozhang · 2024-01-19T03:23:57Z

src/core.rs

+            }
+        }
+
+        assert!(resolved_call.len() == 1);


yihozhang · 2024-01-19T03:25:12Z

src/core.rs

+        ))
+    }
+}
+


I added a MappedExpr type (with brief documentation) as a shorthand

saulshanabrook · 2024-01-19T17:32:45Z

src/serialize.rs

-
-    /// Returns true if the name is in the form v{digits}__
-    /// like v78___
-    ///
-    /// Checks for pattern created by Desugar.get_fresh
-    fn is_temp_name(&self, name: String) -> bool {
-        let number_underscores = self.desugar.number_underscores;
-        let res = name.starts_with('v')
-            && name.ends_with("_".repeat(number_underscores).as_str())
-            && name[1..name.len() - number_underscores]
-                .parse::<u32>()
-                .is_ok();
-        res
-    }


Just checking to see if temporary functions are filtered out in another way now for serialization.

If they are included, it makes the visualization rather noisy!

tests/files.rs

ezrosent · 2024-01-20T20:00:17Z

src/ast/mod.rs

-                }
-                var
-            });
+    pub fn map_def_use(&self, fvar: &mut impl FnMut(&Leaf, bool) -> Leaf) -> Self {


If we can delete them then I think we should :)

Up to you though

src/core.rs

src/lib.rs

ezrosent · 2024-01-20T21:32:49Z

src/lib.rs

-                self.run_command(processed, should_run)?;
+            for processed in self.process_command(command)? {
+                if self.run_mode.show_egglog() {
+                    // This is a hack: we have to do this because process_commands are indeed


Having trouble following what's happening here.

My read here is that show_egglog() is essentially a "dry run" mode where the only commands we want to run are the ones that we need to make everything well-scoped (Push/Pop).

But then how does that relate to typechecking a function twice? When would we do this twice? And aren't we still typechecking in process_command? ?

You are right. I was specifically talking about "typechecking a function twice" because, in our test, we have code like

(sort Expr) (push) (function f () Expr) (pop) (push) (function f () Expr) (pop)

It is only when (push) and (pop) are run that this does not lead to an error. But I agree that "typechecking twice" is not fundamental to the issue here. I'll update the comment to reflect this.

src/main.rs

ezrosent · 2024-01-20T21:35:56Z

src/sort/mod.rs

+}
+
+impl PrimitiveLike for ValueEq {
+    fn name(&self) -> Symbol {


optional:

A while back, I noticed some performance problems related to calling "string constant".into() all over the place. It was faster / cheaper to have a lazy_static! that cached the intern value and then loading that here.

oflatt

Ready to merge! Amazing work 🦸

src/actions.rs

yihozhang and others added 30 commits November 21, 2023 00:14

make type checker happy by commenting out NormFact related functions

a9fcfbb

WIP

bc17d71

type checks after refactoring

b0ebdde

WIP

18bfcb8

typechecks again

0e4c411

incremental progress

cccc9be

some hard issues about typechecking actions

0631e74

type checks

d266085

fix a type error

258bff7

This is not type checking, since action ompilation is still WIP.

bf5dad8

However, the type constraint part is almost done

implemented resolution for exprs, facts, and actions

0df17c4

typechecks again

05d3291

Refactor NormActions to be generic; able to write the to_core_rule fu…

2d03e92

…nction generically.

adapt canonicalize, which eliminates several todos

a5cd649

refactor action checker to use CoreAction

1aab1da

1. Remove compile_expr, use compile_actions instead.

785b776

2. expr_to_norm_actions -> expr.to_norm_actions

math.egg, herbie.egg, knapsack.egg are working

e420230

Some remaining work in main.rs, test.rs, to_resolved, and others.

more tests are running

4573be0

resolve warnings

eb52ade

* Fix several clippy warnings (except for complex types)

168d4ec

* Change the specification of accept * Use indexmaps more to enforce a consistent variable ordering for Load::Subst * Refactor `compile_exprs`

tweak

a735990

big rename of expr to genericexpr

f82c270

fix eval_expr again

eb90b27

Rule -> GenericRule

28d41f2

a bunch of renaming to get rid of ugly unresolved names

95a856b

more renaming

8f4dc1b

fix up nits, allow type complexity

cb8fb1f

Merge pull request #322 from oflatt/oflatt-type-annotation

30426fb

Fix clippy errors and extraction in yihong's typechecking branch

Merge branch 'yihozhang-type-annotation' of github.com:egraphs-good/e…

05b938e

…gglog into yihozhang-type-annotation

names made specialized

eaaccf9

yihozhang added 5 commits December 18, 2023 21:09

reorganize core a bit

da36f7b

fix import errors

4552b30

Added back resugaring, which can be called with --show desugared-egglog

0e15b37

* Added back resugaring, which can be called with --show desugared-eg…

1464c19

…glog * Remove should_run flag

Merge branch 'yihozhang-type-annotation' of github.com:egraphs-good/e…

ec20df5

…gglog into yihozhang-type-annotation

oflatt reviewed Jan 3, 2024

View reviewed changes

src/ast/desugar.rs Outdated Show resolved Hide resolved

yihozhang added 2 commits January 6, 2024 19:19

Move type_info out of Desugar

9cd376f

rename CoreAction/NormAction to GenericCoreAction/CoreAction

e49d54a

oflatt mentioned this pull request Jan 13, 2024

Typechecker tests for variable shadowing #330

Open

ezrosent requested changes Jan 15, 2024

View reviewed changes

Address some easy comments

ed0d7f1

yihozhang commented Jan 19, 2024

View reviewed changes

yihozhang requested a review from ezrosent January 19, 2024 03:47

saulshanabrook reviewed Jan 19, 2024

View reviewed changes

ezrosent approved these changes Jan 20, 2024

View reviewed changes

Address @ezrosent's comments

1ee1cab

yihozhang requested a review from oflatt January 23, 2024 19:31

oflatt approved these changes Jan 24, 2024

View reviewed changes

src/actions.rs Show resolved Hide resolved

yihozhang added 3 commits January 24, 2024 13:59

add a TODO on using only the order and remove the types

a49ded1

Add tests for #80, #113, and #196

f92665c

fix doc

4a2b3dc

yihozhang merged commit 3b83e39 into main Jan 25, 2024
6 checks passed

This was referenced Jan 25, 2024

Query atoms are order-sensitive #196

Closed

Created map type is always first defined map type #113

Closed

saulshanabrook deleted the yihozhang-type-annotation branch January 26, 2024 01:29

ConnorBaker mentioned this pull request Jan 26, 2024

egglog: unstable-2023-09-12 -> 0-unstable-2024-01-26 NixOS/nixpkgs#283906

Merged

13 tasks

This was referenced Feb 7, 2024

Upgrade all dependencies egraphs-good/egglog-python#114

Merged

Remove terms file #342

Merged

This was referenced Mar 28, 2024

Query bug in container instantiation with unrelated containers #308

Closed

Type-annotated AST #211

Closed

		@@ -0,0 +1,755 @@
		/// This file implements the core IR of the language, which is called CoreRule.

Adding type annotations to intermediate representation #320

Adding type annotations to intermediate representation #320

Conversation

yihozhang commented Dec 8, 2023 • edited Loading

ezrosent 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

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

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

yihozhang 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

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

oflatt left a comment

Choose a reason for hiding this comment

yihozhang commented Dec 8, 2023 •

edited

Loading