From 8b7d044ef1849a70ac6f112ac3bafb643a913973 Mon Sep 17 00:00:00 2001
From: Tyler K <tyly04@gmail.com>
Date: Wed, 17 Jul 2024 08:14:31 -0700
Subject: [PATCH] Add disabling support

---
 tool/src/demo_gen/mod.rs      | 10 ++++++++++
 tool/src/demo_gen/terminus.rs | 27 +++++++++++++++++++++------
 tool/src/lib.rs               |  2 ++
 3 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/tool/src/demo_gen/mod.rs b/tool/src/demo_gen/mod.rs
index 16da26e2b..0e9ad5c40 100644
--- a/tool/src/demo_gen/mod.rs
+++ b/tool/src/demo_gen/mod.rs
@@ -71,7 +71,17 @@ impl<'tcx> WebDemoGenerationContext<'tcx> {
 
             {
                 let type_name = self.formatter.fmt_type_name(id);
+                
+                let ty = self.tcx.resolve_type(id);
+                if ty.attrs().disable {
+                    continue;
+                }
+
                 for method in methods {
+                    if method.attrs.disable {
+                        continue;
+                    }
+
                     let val = RenderTerminusContext::evaluate_terminus(
                         self,
                         type_name.to_string(),
diff --git a/tool/src/demo_gen/terminus.rs b/tool/src/demo_gen/terminus.rs
index 7faeb0f50..378d430a0 100644
--- a/tool/src/demo_gen/terminus.rs
+++ b/tool/src/demo_gen/terminus.rs
@@ -216,29 +216,41 @@ impl<'a, 'tcx> RenderTerminusContext<'a, 'tcx> {
                         self.ctx.formatter.fmt_primitive_list_type(*p).to_string()
                     }
                     Type::Slice(hir::Slice::Strs(..)) => "Array<String>".to_string(),
-                    Type::Enum(e) => self
+                    Type::Enum(e) => {
+                        let type_name = self
                         .ctx
                         .formatter
                         .fmt_type_name(e.tcx_id.into())
-                        .to_string(),
+                        .to_string();
+                        
+                        if e.resolve(&self.ctx.tcx).attrs.disable {
+                            self.ctx.errors.push_error(format!("Found usage of disabled type {type_name}"))
+                        }
+
+                        type_name
+                    },
                     _ => unreachable!("Unknown primitive type {:?}", param_type),
                 };
 
                 out_param(type_name);
             }
             Type::Opaque(o) => {
-                let attrs = &o.resolve(self.ctx.tcx).attrs.demo_attrs;
-
-                // We need to find a constructor that we can call.
-                // Piggybacking off of the #[diplomat::attr(constructor)] macro for now as well as test attributes in attrs.rs
                 let op = o.resolve(self.ctx.tcx);
                 let type_name = self.ctx.formatter.fmt_type_name(o.tcx_id.into());
 
+                let all_attrs = &o.resolve(self.ctx.tcx).attrs;
+                if all_attrs.disable {
+                    self.ctx.errors.push_error(format!("Found usage of disabled type {type_name}"))
+                }
+                let attrs = &all_attrs.demo_attrs;
+
                 if attrs.external {
                     out_param(type_name.into());
                     return;
                 }
 
+                // We need to find a constructor that we can call.
+                // Piggybacking off of the #[diplomat::attr(constructor)] macro for now as well as test attributes in attrs.rs
                 let mut usable_constructor = false;
 
                 for method in op.methods.iter() {
@@ -283,6 +295,9 @@ impl<'a, 'tcx> RenderTerminusContext<'a, 'tcx> {
                 let st = s.resolve(self.ctx.tcx);
 
                 let type_name = self.ctx.formatter.fmt_type_name(s.tcx_id.into());
+                if st.attrs.disable {
+                    self.ctx.errors.push_error(format!("Found usage of disabled type {type_name}"))
+                }
 
                 self.terminus_info
                     .imports
diff --git a/tool/src/lib.rs b/tool/src/lib.rs
index 6af84135e..b23ae05f5 100644
--- a/tool/src/lib.rs
+++ b/tool/src/lib.rs
@@ -113,6 +113,8 @@ pub fn gen(
         "demo-gen" => {
             let mut attr_validator = hir::BasicAttributeValidator::new("demo-gen");
 
+            attr_validator.support.disabling = true;
+
             // For finding default constructors of opaques:
             attr_validator.support.constructors = true;
             attr_validator.support.fallible_constructors = true;