Fix target-aware prepare path and add regression test

Agent-Logs-Url: https://github.com/microsoft/regorus/sessions/f964122b-5af3-41a7-bc51-d87b7271a455

Co-authored-by: anakrish <35780660+anakrish@users.noreply.github.com>

src/engine.rs

@@ -1125,6 +1125,10 @@ impl Engine {
         limits::enforce_memory_limit().map_err(|err| anyhow!(err))?;
 
         self.interpreter.set_traces(enable_tracing);
+        #[cfg(feature = "azure_policy")]
+        let was_prepared = self.prepared;
+        #[cfg(not(feature = "azure_policy"))]
+        let _ = for_target;
 
         // if the data/policies have changed or the interpreter has never been prepared
         if !self.prepared {
@@ -1158,25 +1162,22 @@ impl Engine {
             // Set schedule after hoisting completes
             self.interpreter.set_schedule(Some(schedule));
 
-            #[cfg(feature = "azure_policy")]
-            if for_target {
-                // Resolve and validate target specifications across all modules
-                crate::interpreter::target::resolve::resolve_and_apply_target(
-                    &mut self.interpreter,
-                )?;
-                // Infer resource types
-                crate::interpreter::target::infer::infer_resource_type(&mut self.interpreter)?;
-            }
-
-            if !for_target {
-                // Check if any module specifies a target and warn if so
-                #[cfg(feature = "azure_policy")]
-                self.warn_if_targets_present();
-            }
-
             self.prepared = true;
         }
 
+        #[cfg(feature = "azure_policy")]
+        if for_target {
+            // Resolve and validate target specifications across all modules.
+            // This must run for target-aware compilation even if generic prepare()
+            // was already called.
+            crate::interpreter::target::resolve::resolve_and_apply_target(&mut self.interpreter)?;
+            // Infer resource types
+            crate::interpreter::target::infer::infer_resource_type(&mut self.interpreter)?;
+        } else if !was_prepared {
+            // Check if any module specifies a target and warn if so.
+            self.warn_if_targets_present();
+        }
+
         Ok(())
     }
 

tests/engine/mod.rs

@@ -142,6 +142,46 @@ fn prepare_then_clone_without_initial_eval() -> Result<()> {
     Ok(())
 }
 
+#[test]
+#[cfg(feature = "azure_policy")]
+#[cfg_attr(docsrs, doc(cfg(feature = "azure_policy")))]
+fn prepare_then_compile_for_target() -> Result<()> {
+    if !registry::targets::contains("target.tests.sample_test_target") {
+        let target = Target::from_json_str(include_str!(
+            "../interpreter/cases/target/definitions/sample_target.json"
+        ))?;
+        registry::targets::register(Rc::new(target))?;
+    }
+
+    let mut engine = Engine::new();
+    engine.add_policy(
+        "test.rego".to_string(),
+        r#"package test
+           import rego.v1
+           __target__ := "target.tests.sample_test_target"
+
+           default allow := false
+
+           allow if {
+             input.type == "test_resource"
+           }
+        "#
+        .to_string(),
+    )?;
+
+    engine.prepare()?;
+    let compiled = engine.compile_for_target()?;
+    let info = compiled.get_policy_info()?;
+    assert_eq!(info.target_name.as_deref(), Some("target.tests.sample_test_target"));
+
+    let result = compiled.eval_with_input(Value::from_json_str(
+        r#"{"name":"resource-1","type":"test_resource"}"#,
+    )?)?;
+    assert_eq!(result, Value::from(true));
+
+    Ok(())
+}
+
 #[test]
 #[cfg(feature = "azure_policy")]
 #[cfg_attr(docsrs, doc(cfg(feature = "azure_policy")))]