From 7c3c919fe89bd8bc29a18f39a402f94ef12f4446 Mon Sep 17 00:00:00 2001
From: LunaAmora <luna.mberry@gmail.com>
Date: Fri, 7 Jun 2024 11:13:46 -0300
Subject: [PATCH] Limit adt tag defs be generated or builtin

---
 src/fun/mod.rs                                     |  2 +-
 src/fun/parser.rs                                  |  4 +++-
 src/fun/transform/definition_pruning.rs            | 13 ++++++++-----
 src/fun/transform/encode_adts.rs                   |  6 +++---
 .../scott_triggers_unused__test.bend.snap          | 14 +++++++++-----
 5 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/src/fun/mod.rs b/src/fun/mod.rs
index 6925039f..9f0dc1d6 100644
--- a/src/fun/mod.rs
+++ b/src/fun/mod.rs
@@ -88,7 +88,7 @@ pub enum Source {
 pub struct HvmDefinition {
   pub name: Name,
   pub body: hvm::ast::Net,
-  pub builtin: bool,
+  pub source: Source,
 }
 
 /// A pattern matching rule of a definition.
diff --git a/src/fun/parser.rs b/src/fun/parser.rs
index 0467e2cf..181a6550 100644
--- a/src/fun/parser.rs
+++ b/src/fun/parser.rs
@@ -252,7 +252,9 @@ impl<'a> TermParser<'a> {
     let mut p = hvm::ast::CoreParser::new(&self.input[*self.index()..]);
     let body = p.parse_net()?;
     *self.index() = ini_idx + *p.index();
-    let def = HvmDefinition { name: name.clone(), body, builtin };
+    let end_idx = *self.index();
+    let source = if builtin { Source::Builtin } else { Source::Local(ini_idx..end_idx) };
+    let def = HvmDefinition { name: name.clone(), body, source };
     Ok(def)
   }
 
diff --git a/src/fun/transform/definition_pruning.rs b/src/fun/transform/definition_pruning.rs
index a8c62d31..0e9ba616 100644
--- a/src/fun/transform/definition_pruning.rs
+++ b/src/fun/transform/definition_pruning.rs
@@ -1,6 +1,6 @@
 use crate::{
   diagnostics::WarningType,
-  fun::{Book, Ctx, Name, Term},
+  fun::{Book, Ctx, Name, Source, Term},
   maybe_grow,
 };
 use hvm::ast::{Net, Tree};
@@ -43,7 +43,7 @@ impl Ctx<'_> {
       }
     }
     for def in self.book.hvm_defs.values() {
-      if !def.builtin && !(used.get(&def.name) == Some(&Used::Main)) {
+      if !def.source.is_builtin() && !(used.get(&def.name) == Some(&Used::Main)) {
         used.insert(def.name.clone(), Used::NonBuiltin);
         self.book.find_used_definitions_from_hvm_net(&def.body, Used::NonBuiltin, &mut used);
       }
@@ -60,8 +60,11 @@ impl Ctx<'_> {
     }
 
     // Remove unused definitions.
-    let names = self.book.defs.keys().cloned().chain(self.book.hvm_defs.keys().cloned()).collect::<Vec<_>>();
-    for def in names {
+    let defs = self.book.defs.iter().map(|(nam, def)| (nam.clone(), def.source.clone()));
+    let hvm_defs = self.book.hvm_defs.iter().map(|(nam, def)| (nam.clone(), def.source.clone()));
+    let names = defs.chain(hvm_defs).collect::<Vec<_>>();
+
+    for (def, src) in names {
       if let Some(use_) = used.get(&def) {
         match use_ {
           Used::Main => {
@@ -72,7 +75,7 @@ impl Ctx<'_> {
             // Prune if `prune_all`, otherwise show a warning.
             if prune_all {
               rm_def(self.book, &def);
-            } else if !def.is_generated() {
+            } else if !def.is_generated() && !matches!(src, Source::Generated) {
               self.info.add_rule_warning("Definition is unused.", WarningType::UnusedDefinition, def);
             }
           }
diff --git a/src/fun/transform/encode_adts.rs b/src/fun/transform/encode_adts.rs
index c5c9c09e..d937c0c1 100644
--- a/src/fun/transform/encode_adts.rs
+++ b/src/fun/transform/encode_adts.rs
@@ -18,7 +18,7 @@ impl Book {
           AdtEncoding::NumScott => {
             let tag = make_tag(adt_name == ctr_name, ctr_name);
             let body = encode_ctr_num_scott(fields.iter().map(|f| &f.nam), &tag);
-            let tag_def = make_tag_def(ctr_idx, &tag, adt);
+            let tag_def = make_tag_def(ctr_idx, &tag, adt.source.is_builtin());
             tags.push((tag, tag_def));
             body
           }
@@ -65,7 +65,7 @@ fn encode_ctr_num_scott<'a>(ctr_args: impl DoubleEndedIterator<Item = &'a Name>
   Term::rfold_lams(term, ctr_args.cloned().map(Some))
 }
 
-fn make_tag_def(ctr_idx: usize, tag: &Name, adt: &crate::fun::Adt) -> Definition {
+fn make_tag_def(ctr_idx: usize, tag: &Name, builtin: bool) -> Definition {
   let tag_rule = vec![Rule { pats: vec![], body: Term::Num { val: Num::U24(ctr_idx as u32) } }];
-  Definition::new(tag.clone(), tag_rule, adt.source.clone())
+  Definition::new_gen(tag.clone(), tag_rule, builtin)
 }
diff --git a/tests/snapshots/scott_triggers_unused__test.bend.snap b/tests/snapshots/scott_triggers_unused__test.bend.snap
index e94d7ece..e7221e2c 100644
--- a/tests/snapshots/scott_triggers_unused__test.bend.snap
+++ b/tests/snapshots/scott_triggers_unused__test.bend.snap
@@ -2,8 +2,12 @@
 source: tests/golden_tests.rs
 input_file: tests/golden_tests/scott_triggers_unused/test.bend
 ---
-Errors:
-In definition 'bool/f/tag':
-  Definition is unused.
-In definition 'bool/t/tag':
-  Definition is unused.
+@bool/f = ((@bool/f/tag a) a)
+
+@bool/f/tag = 1
+
+@bool/t = ((@bool/t/tag a) a)
+
+@bool/t/tag = 0
+
+@main = (((?((0 (* 1)) a) a) b) b)