fold_file -define

Summary: As title

Reviewed By: perehonchuk

Differential Revision: D49821862

fbshipit-source-id: be9376b6e73fd92bc519c5d0fa168d41ee03cd04

crates/hir/src/fold.rs

@@ -30,6 +30,7 @@ use crate::FunType;
 use crate::HirIdx;
 use crate::InFile;
 use crate::ListType;
+use crate::PPDirective;
 use crate::Pat;
 use crate::PatId;
 use crate::RecordFieldBody;
@@ -90,6 +91,18 @@ pub fn fold_file<'a, T>(
                     callback(acc, ctx)
                 })
             }
+            FormIdx::PPDirective(idx) => {
+                if let PPDirective::Define(define_id) = &form_list[idx] {
+                    sema.fold_define(
+                        form_idx,
+                        InFile::new(file_id, *define_id),
+                        r,
+                        &mut |acc, ctx| callback(acc, ctx),
+                    )
+                } else {
+                    r
+                }
+            }
             _ => {
                 // Will have to do some time?
                 r
@@ -1667,4 +1680,13 @@ bar() ->
                "#;
         count_atom_foo(fixture_str, 2);
     }
+
+    #[test]
+    fn traverse_define() {
+        let fixture_str = r#"
+               -module(foo).
+               -define(FOO(X), foo(X,fo~o)).
+               "#;
+        count_atom_foo(fixture_str, 2);
+    }
 }

crates/hir/src/sema.rs

@@ -56,6 +56,7 @@ use crate::CallbackId;
 use crate::Clause;
 use crate::CompileOptionId;
 use crate::DefMap;
+use crate::DefineId;
 use crate::Expr;
 use crate::ExprId;
 use crate::File;
@@ -673,6 +674,27 @@ impl<'db> Semantic<'db> {
         )
     }
 
+    pub fn fold_define<'a, T>(
+        &self,
+        form_id: FormIdx,
+        attribute_id: InFile<DefineId>,
+        initial: T,
+        callback: AnyCallBack<'a, T>,
+    ) -> T {
+        if let Some(body) = self.db.define_body(attribute_id) {
+            FoldCtx::fold_expr(
+                &body.body,
+                Strategy::TopDown,
+                form_id,
+                body.expr,
+                initial,
+                callback,
+            )
+        } else {
+            initial
+        }
+    }
+
     pub fn bound_vars_in_pattern_diagnostic(
         &self,
         file_id: FileId,