Add support for instrumentation rules

momvart · momvart · commit 68b157e5c0d8 · 2024-12-29T02:11:21.000Z
diff --git a/compiler/Cargo.toml b/compiler/Cargo.toml
@@ -29,6 +29,7 @@ glob = "0.3.1"
 config = { workspace = true }
 serde = { workspace = true }
 ron = "0.8"
+regex-lite = "0.1"
 bitflags = "2.6.0"
 tracing-subscriber = { workspace = true }
 
diff --git a/compiler/src/config.rs b/compiler/src/config.rs
@@ -15,6 +15,8 @@ pub(crate) struct LeafCompilerConfig {
     pub codegen_all_mir: bool,
     #[serde(default = "default_marker_cfg_name")]
     pub marker_cfg_name: String,
+    #[serde(default)]
+    pub rules: InstrumentationRules,
 }
 
 fn default_override_sysroot() -> bool {
@@ -92,6 +94,83 @@ fn default_runtime_shim_crate_name() -> String {
     "leaf".to_string()
 }
 
+/* NOTE: How is serde's structure is defined?
+ * We want to make the rules easy and intuitive to define in TOML.
+ * - The default enum representation in serde uses the variant name as the key.
+ * - The untagged representation selects the variant based on unique fields matched.
+ * We mostly utilize these two and flattening.
+ * For example, a `LogicFormula` can be represented as any of the following:
+ * ```toml
+ * [[f]]
+ * crate = { is_external = true }
+ * [[f]]
+ * not = { crate = { name = "std" } }
+ * [[f]]
+ * any = [{ crate = { name = "std" } }, { crate = { name = "core" } }]
+ * [[f]]
+ * all = [{ crate = { is_external = true } }, { crate = { name = "core" } }]
+ * ```
+ */
+
+#[derive(Debug, Clone, Deserialize, Default)]
+pub(crate) struct InstrumentationRules {
+    #[serde(default)]
+    pub(crate) include: Vec<EntityFilter>,
+    #[serde(default)]
+    pub(crate) exclude: Vec<EntityFilter>,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+#[serde(untagged)]
+pub(crate) enum EntityFilter {
+    WholeBody(LogicFormula<EntityLocationFilter>),
+}
+
+#[derive(Debug, Clone, Deserialize)]
+#[serde(untagged)]
+pub(crate) enum LogicFormula<T> {
+    Not(NotFormula<T>),
+    Any(AnyFormula<T>),
+    All(AllFormula<T>),
+    Atom(T),
+    // NOTE: This variant helps with parsing empty tables by preventing the infinite search over the name of fields.
+    None {},
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub(crate) struct NotFormula<T> {
+    #[serde(rename = "not")]
+    pub(crate) of: Box<LogicFormula<T>>,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub(crate) struct AnyFormula<T> {
+    #[serde(rename = "any")]
+    pub(crate) of: Vec<LogicFormula<T>>,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub(crate) struct AllFormula<T> {
+    #[serde(rename = "all")]
+    pub(crate) of: Vec<LogicFormula<T>>,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub(crate) enum EntityLocationFilter {
+    #[serde(alias = "def_path")]
+    DefPathMatch(String),
+    Crate(CrateFilter),
+}
+
+#[derive(Debug, Clone, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub(crate) enum CrateFilter {
+    #[serde(alias = "is_external")]
+    Externality(bool),
+    Name(String),
+}
+
 const CONFIG_FILENAME: &str = "leafc_config";
 
 pub(super) fn load_config() -> LeafCompilerConfig {
diff --git a/compiler/src/main.rs b/compiler/src/main.rs
@@ -203,7 +203,7 @@ mod driver_callbacks {
             <LeafToolAdder>,
             <TypeExporter>,
             nctfe_pass,
-            Instrumentor::new(true, None /* FIXME */),
+            Instrumentor::new(true, None /* FIXME */, config.rules.clone()),
         );
 
         if config.codegen_all_mir {
diff --git a/compiler/src/passes/instr/decision.rs b/compiler/src/passes/instr/decision.rs
@@ -7,31 +7,44 @@ use rustc_middle::{
 };
 use rustc_span::Symbol;
 
-use common::{log_info, log_warn};
+use core::ops::Deref;
 
-use crate::utils::mir::TyCtxtExt;
+use common::{log_debug, log_info, log_warn};
 
-pub(super) const TAG_INSTR_DECISION: &str = concatcp!(super::TAG_INSTRUMENTATION, "::skip");
+use crate::{config::InstrumentationRules, utils::mir::TyCtxtExt};
+
+pub(super) const TAG_INSTR_DECISION: &str = concatcp!(super::TAG_INSTRUMENTATION, "::decision");
 
 const TOOL_NAME: &str = crate::constants::TOOL_LEAF;
 const ATTR_NAME: &str = "instrument";
+pub(super) const KEY_RULES: &str = "instr_rules";
+const KEY_BAKED_RULES: &str = "instr_rules_baked";
 
-pub(super) fn should_instrument<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool {
+pub(super) fn should_instrument<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    body: &mut Body<'tcx>,
+    storage: &mut dyn Storage,
+) -> bool {
     let def_id = body.source.def_id();
 
     if !decide_instance_kind(&body.source.instance) {
         return false;
     }
 
-    if let Some((explicit, item)) = opt_instrument_attr_inheritable(tcx, def_id) {
-        log_info!(
+    let rules = storage.get_or_insert_with_acc(KEY_BAKED_RULES.to_owned(), |storage| {
+        let rules = storage.get_or_default::<InstrumentationRules>(KEY_RULES.to_owned());
+        rules.to_baked()
+    });
+
+    if let Some((decision, item)) = find_inheritable_first_filtered(tcx, def_id, rules.deref()) {
+        log_debug!(
             target: TAG_INSTR_DECISION,
-            "Found explicit instrumentation attribute for {:?} on {:?} with value: {}",
+            "Found a rule for instrumentation of {:?} on {:?} with decision: {}",
             def_id,
             item,
-            explicit
+            decision
         );
-        return explicit;
+        return decision;
     }
 
     if is_lang_start_item(tcx, def_id) {
@@ -81,16 +94,27 @@ fn decide_instance_kind(kind: &InstanceKind) -> bool {
     }
 }
 
-/// Returns the value of the `instrument` attribute if it is placed on the item or one of its ancestors.
-fn opt_instrument_attr_inheritable<'tcx>(
+fn find_inheritable_first_filtered<'tcx>(
     tcx: TyCtxt<'tcx>,
     def_id: DefId,
+    rules: &dyn Fn(TyCtxt<'tcx>, DefId) -> Option<bool>,
 ) -> Option<(bool, DefId)> {
     let mut current = def_id;
     loop {
-        let attr = opt_instrument_attr(tcx, current);
-        if attr.is_some() {
-            return attr.map(|v| (v, current));
+        // Attributes take precedence over filters.
+        if let Some(explicit) = opt_instrument_attr(tcx, current) {
+            log_info!(
+                target: TAG_INSTR_DECISION,
+                "Found explicit instrumentation attribute for {:?} on {:?} with value: {}",
+                def_id,
+                current,
+                explicit
+            );
+            return Some((explicit, current));
+        }
+
+        if let Some(include) = rules(tcx, current) {
+            return Some((include, current));
         }
 
         let parent = tcx.opt_parent(current);
@@ -817,3 +841,91 @@ mod intrinsics {
     }
 }
 pub(super) use intrinsics::{decide_intrinsic_call, AtomicIntrinsicKind, IntrinsicDecision};
+
+use super::{Storage, StorageExt};
+
+mod rules {
+    use super::*;
+
+    use crate::config::{
+        AllFormula, AnyFormula, CrateFilter, EntityFilter, EntityLocationFilter,
+        InstrumentationRules, LogicFormula, NotFormula,
+    };
+
+    type Predicate = Box<dyn Fn(TyCtxt<'_>, DefId) -> bool>;
+
+    trait ToPredicate {
+        fn to_predicate<'tcx>(&self) -> Predicate;
+    }
+
+    impl InstrumentationRules {
+        pub(super) fn to_baked<'tcx>(&self) -> impl Fn(TyCtxt<'_>, DefId) -> Option<bool> {
+            fn to_predicate(filters: &[EntityFilter]) -> Predicate {
+                let preds = filters
+                    .iter()
+                    .filter_map(|f| match &f {
+                        EntityFilter::WholeBody(formula) => Some(formula),
+                    })
+                    .map(|f| f.to_predicate())
+                    .collect::<Vec<_>>();
+                Box::new(move |tcx, def_id| preds.iter().any(|p| p(tcx, def_id)))
+            }
+
+            let include = to_predicate(&self.include);
+            let exclude = to_predicate(&self.exclude);
+
+            move |tcx, def_id| {
+                if include(tcx, def_id) {
+                    Some(true)
+                } else if exclude(tcx, def_id) {
+                    Some(false)
+                } else {
+                    None
+                }
+            }
+        }
+    }
+
+    impl<T: ToPredicate> ToPredicate for LogicFormula<T> {
+        fn to_predicate<'tcx>(&self) -> Predicate {
+            match self {
+                LogicFormula::None {} => Box::new(|_, _| false),
+                LogicFormula::Atom(f) => f.to_predicate(),
+                LogicFormula::Not(NotFormula { of }) => {
+                    let pred = of.to_predicate();
+                    Box::new(move |tcx, def_id| !pred(tcx, def_id))
+                }
+                LogicFormula::Any(AnyFormula { of }) => {
+                    let preds = of.iter().map(|f| f.to_predicate()).collect::<Vec<_>>();
+                    Box::new(move |tcx, def_id| preds.iter().any(|p| p(tcx, def_id)))
+                }
+                LogicFormula::All(AllFormula { of }) => {
+                    let preds = of.iter().map(|f| f.to_predicate()).collect::<Vec<_>>();
+                    Box::new(move |tcx, def_id| preds.iter().all(|p| p(tcx, def_id)))
+                }
+            }
+        }
+    }
+
+    impl ToPredicate for EntityLocationFilter {
+        fn to_predicate<'tcx>(&self) -> Predicate {
+            match self {
+                EntityLocationFilter::Crate(crate_filter) => match crate_filter.clone() {
+                    CrateFilter::Externality(is_external) => {
+                        Box::new(move |_, def_id| def_id.is_local() != is_external)
+                    }
+                    CrateFilter::Name(name) => {
+                        Box::new(move |tcx, def_id| tcx.crate_name(def_id.krate).as_str() == name)
+                    }
+                },
+                EntityLocationFilter::DefPathMatch(pattern) => {
+                    let regex = regex_lite::Regex::new(&pattern).expect("Invalid regex pattern");
+                    Box::new(move |tcx, def_id| {
+                        let def_path = tcx.def_path_str(def_id);
+                        regex.is_match(&def_path)
+                    })
+                }
+            }
+        }
+    }
+}
diff --git a/compiler/src/passes/instr/mod.rs b/compiler/src/passes/instr/mod.rs
@@ -19,6 +19,7 @@ use common::{log_debug, log_info, log_warn};
 use std::{collections::HashSet, num::NonZeroUsize, sync::atomic};
 
 use crate::{
+    config::InstrumentationRules,
     mir_transform::{self, BodyInstrumentationUnit, JumpTargetModifier},
     passes::{instr::call::context::PriItems, StorageExt},
     utils::mir::TyCtxtExt,
@@ -50,13 +51,19 @@ const KEY_TOTAL_COUNT: &str = "total_body_count";
 pub(crate) struct Instrumentor {
     enabled: bool,
     total_body_count: Option<NonZeroUsize>,
+    rules: Option<InstrumentationRules>,
 }
 
 impl Instrumentor {
-    pub(crate) fn new(enabled: bool, total_body_count: Option<NonZeroUsize>) -> Self {
+    pub(crate) fn new(
+        enabled: bool,
+        total_body_count: Option<NonZeroUsize>,
+        filters: InstrumentationRules,
+    ) -> Self {
         Self {
             enabled,
             total_body_count,
+            rules: Some(filters),
         }
     }
 }
@@ -76,6 +83,9 @@ impl CompilationPass for Instrumentor {
         // As early as possible, we use transform_ast to set the enabled flag.
         storage.get_or_insert_with(KEY_ENABLED.to_owned(), || self.enabled);
         storage.get_or_insert_with(KEY_TOTAL_COUNT.to_owned(), || self.total_body_count);
+        storage.get_or_insert_with(decision::KEY_RULES.to_owned(), || {
+            self.rules.take().unwrap()
+        });
         rustc_driver::Compilation::Continue
     }
 
@@ -97,7 +107,7 @@ fn transform<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, storage: &mut dyn S
 
     let def_id = body.source.def_id();
 
-    if !decision::should_instrument(tcx, body) {
+    if !decision::should_instrument(tcx, body, storage) {
         log_info!(
             target: decision::TAG_INSTR_DECISION,
             "Skipping instrumentation for {:#?}",
diff --git a/compiler/src/passes/mod.rs b/compiler/src/passes/mod.rs
