Add --auto-defines and --suggest-defines macro census flags

src/main/scala/io/appthreat/atom/Atom.scala

@@ -416,6 +416,25 @@ object Atom:
               case config: AtomConfig => config.withAppendedFrontendArg("defines", x)
               case _                  => c
         )
+    opt[Unit]("auto-defines")
+        .text(
+          "Run a macro census first and define the build-option macros (CONFIG_*, ENABLE_*, template-declared) that hide #if code. Opt-in: it changes what is analysed. (C/C++ only)"
+        )
+        .action((_, c) =>
+            c match
+              case config: AtomConfig => config.withFrontendArg("auto-defines", "true")
+              case _                  => c
+        )
+    opt[String]("suggest-defines")
+        .valueName("<file>")
+        .text(
+          "Write the macro census to <file>.json and a reviewable --macro-files header to <file>.h, then exit; with --auto-defines, write it and continue. (C/C++ only)"
+        )
+        .action((x, c) =>
+            c match
+              case config: AtomConfig => config.withFrontendArg("macro-census", x)
+              case _                  => c
+        )
     opt[String]("include-path")
         .unbounded()
         .text("Header include path. Repeatable. (C/C++ only)")
@@ -842,10 +861,27 @@ object Atom:
                 if loaded.frontendArgsKeys then
                   println(FrontendArgsApplier.renderKeys(loaded.language))
                   Right("Displayed frontend-args keys")
-                else run(loaded, loaded.language)
+                else if isCensusOnly(loaded) then
+                  Try(
+                    new C2Cpg().writeMacroCensus(
+                      c2CpgConfig(loaded, loaded.outputAtomFile.pathAsString)
+                    )
+                  ).toEither.left.map(e => s"Macro census failed: ${e.getMessage}")
+                      .map(_ => "Wrote the macro census")
+                else
+                  if hasCensusFlags(loaded) && !CensusLanguages.contains(
+                      loaded.language.toUpperCase
+                    )
+                  then
+                    println(
+                      s"--auto-defines and --suggest-defines apply to C/C++ only; ignored for -l ${loaded.language}"
+                    )
+                  run(loaded, loaded.language)
             case Left(err) => Left(err)
       case Right(_)  => Left("Invalid configuration generated")
       case Left(err) => Left(err)
+    end match
+  end run
 
   private def run(config: AtomConfig, language: String): Either[String, String] =
       for
@@ -1228,9 +1264,12 @@ object Atom:
         .withIncludeTrivialExpressions(false)
         .withIncludePaths(C2ATOM_INCLUDE_PATH)
     val finalConfig = FrontendArgsApplier.applyC(baseConfig, config.frontendArgs)
-    new C2Atom().createCpg(finalConfig)
+    new C2Atom().createCpg(C2Cpg.withCensusDefines(finalConfig))
 
   private def createC2Cpg(config: AtomConfig, outputAtomFile: String): Try[Cpg] =
+      new C2Cpg().createCpgWithOverlays(c2CpgConfig(config, outputAtomFile))
+
+  private def c2CpgConfig(config: AtomConfig, outputAtomFile: String): CConfig =
     val cIgnoreDirEnvVars =
         if config.language.equalsIgnoreCase("CPP") || config.language.equalsIgnoreCase("C++") then
           Seq("CHEN_C_IGNORE_DIRS", "CHEN_CPP_IGNORE_DIRS")
@@ -1249,9 +1288,24 @@ object Atom:
         .withImageLocations(false)
         .withIncludeTrivialExpressions(false)
         .withIncludePaths(C2ATOM_INCLUDE_PATH)
-    val finalConfig = FrontendArgsApplier.applyC(baseConfig, config.frontendArgs)
-    new C2Cpg().createCpgWithOverlays(finalConfig)
-  end createC2Cpg
+    FrontendArgsApplier.applyC(baseConfig, config.frontendArgs)
+  end c2CpgConfig
+
+  /** `--suggest-defines` without `--auto-defines`: the macro census alone, no atom. */
+  /** Every language the C frontends serve, so every one runs the census the same way. */
+  private val CensusLanguages = Set("C", "NEWC", "CPP", "C++", "H", "HPP", "I")
+
+  private def hasCensusFlags(config: AtomConfig): Boolean =
+      config.frontendArgs.contains("auto-defines") || config.frontendArgs.contains("macro-census")
+
+  /** `--suggest-defines` without `--auto-defines`: the census alone, no atom. Read from the
+    * resolved frontend config, so the flags mean exactly what the frontend will see.
+    */
+  private def isCensusOnly(config: AtomConfig): Boolean =
+      hasCensusFlags(config) && CensusLanguages.contains(config.language.toUpperCase) && {
+          val c = c2CpgConfig(config, config.outputAtomFile.pathAsString)
+          c.macroCensusReport.nonEmpty && !c.autoDefines
+      }
 
   private def createJimple2Cpg(config: AtomConfig, outputAtomFile: String): Try[Cpg] =
     val baseConfig = JimpleConfig(android = androidJarPath, fullResolver = true)

src/main/scala/io/appthreat/atom/frontends/FrontendArgsApplier.scala

@@ -104,6 +104,8 @@ object FrontendArgsApplier:
     r = r.withLogProblems(bool(args, "log-problems", r.logProblems))
     r = r.withLogPreprocessor(bool(args, "log-preprocessor", r.logPreprocessor))
     r = r.withPrintIfDefsOnly(bool(args, "print-ifdef-only", r.printIfDefsOnly))
+    r = r.withAutoDefines(bool(args, "auto-defines", r.autoDefines))
+    r = r.withMacroCensusReport(str(args, "macro-census", r.macroCensusReport))
     r = r.withIncludePathsAutoDiscovery(bool(
       args,
       "include-auto-discovery",
@@ -279,6 +281,20 @@ object FrontendArgsApplier:
     KeyDoc("log-problems", "bool", "false", "Log every parse problem.", CLike),
     KeyDoc("log-preprocessor", "bool", "false", "Log every preprocessor statement.", CLike),
     KeyDoc("print-ifdef-only", "bool", "false", "Print preprocessor conditionals and exit.", CLike),
+    KeyDoc(
+      "auto-defines",
+      "bool",
+      "false",
+      "Run the macro census and define its build-option macros before parsing.",
+      CLike
+    ),
+    KeyDoc(
+      "macro-census",
+      "string",
+      "",
+      "Write the macro census to <file>.json and <file>.h.",
+      CLike
+    ),
     KeyDoc("include-auto-discovery", "bool", "false", "Auto-discover system header paths.", CLike),
     KeyDoc("function-bodies", "bool", "true", "Parse function and method bodies.", CLike),
     KeyDoc(

src/main/scala/io/appthreat/atom/package.scala

@@ -85,8 +85,11 @@ package object atom:
       this.cacheFragments = x
       this
 
+    /** Merges into the arguments the first-class flags (`--define`, `--auto-defines`, ...) have
+      * already set: a later `--frontend-args` overrides the keys it names and keeps the rest.
+      */
     def withFrontendArgs(args: Map[String, String]): AtomConfig =
-      this.frontendArgs = args
+      this.frontendArgs = this.frontendArgs ++ args
       this
 
     /** Sets a single frontend argument, replacing any previous value for the key. Used by the

src/test/scala/io/appthreat/atom/MacroCensusCommandTests.scala

@@ -0,0 +1,77 @@
+package io.appthreat.atom
+
+import better.files.File as BFile
+import io.shiftleft.codepropertygraph.cpgloading.{CpgLoader, CpgLoaderConfig}
+import io.shiftleft.semanticcpg.language.*
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.wordspec.AnyWordSpec
+import overflowdb.Config as OdbConfig
+
+/** `--suggest-defines` writes the macro census and stops; `--auto-defines` applies it. Both are
+  * first-class flags that a later `--frontend-args` must not drop.
+  */
+class MacroCensusCommandTests extends AnyWordSpec with Matchers:
+
+  private val source =
+      """#if CONFIG_FEATURE
+        |int feature_fn(int x) {
+        |  x += 1;
+        |  x += 2;
+        |  return x;
+        |}
+        |#endif
+        |int always(int x) { return x; }
+        |""".stripMargin
+
+  private def methodNames(atomFile: BFile): List[String] =
+    val odbConfig = OdbConfig.withDefaults().withStorageLocation(atomFile.pathAsString)
+    val cpg = CpgLoader.loadFromOverflowDb(
+      CpgLoaderConfig().withOverflowConfig(odbConfig).doNotCreateIndexesOnLoad
+    )
+    try cpg.method.name.l
+    finally cpg.close()
+
+  "--suggest-defines" should:
+    "write the census and build no atom, even before a --frontend-args" in:
+      BFile.usingTemporaryDirectory("atom-census") { dir =>
+        val src = dir / "src"
+        src.createDirectory()
+        (src / "a.c").write(source)
+        val atomFile = dir / "app.atom"
+        val base     = dir / "out" / "census"
+        val result = Atom.run(
+          Array(
+            "-l",
+            "c",
+            "-o",
+            atomFile.pathAsString,
+            "--suggest-defines",
+            base.pathAsString,
+            "--frontend-args",
+            "enable-ast-cache=false",
+            src.pathAsString
+          )
+        )
+        result shouldBe Right("Wrote the macro census")
+        (dir / "out" / "census.h").contentAsString should include("#define CONFIG_FEATURE 1")
+        atomFile.exists shouldBe false
+      }
+
+  "--auto-defines" should:
+    "parse the code behind an auto-tier macro" in:
+      BFile.usingTemporaryDirectory("atom-census-auto") { dir =>
+        val src = dir / "src"
+        src.createDirectory()
+        (src / "a.c").write(source)
+        val plain = dir / "plain.atom"
+        val auto  = dir / "auto.atom"
+        def build(out: BFile, extra: String*) = Atom.run(
+          (Seq("-l", "c", "-o", out.pathAsString, "--frontend-args", "enable-ast-cache=false") ++
+              extra ++ Seq(src.pathAsString)).toArray
+        )
+        build(plain).isRight shouldBe true
+        build(auto, "--auto-defines").isRight shouldBe true
+        methodNames(plain) should not contain "feature_fn"
+        methodNames(auto) should contain("feature_fn")
+      }
+end MacroCensusCommandTests