From a1ed3d8ffb8e86c60cc2fc41adbf1855c4a2e0af Mon Sep 17 00:00:00 2001
From: Nathan Froyd <froydnj@gmail.com>
Date: Wed, 30 Aug 2023 14:15:33 -0400
Subject: [PATCH] [rust] fix accessing optional locations

---
 rust/yarp/build.rs   |  3 ++-
 rust/yarp/src/lib.rs | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/rust/yarp/build.rs b/rust/yarp/build.rs
index df90a06f211..a623926f512 100644
--- a/rust/yarp/build.rs
+++ b/rust/yarp/build.rs
@@ -246,7 +246,8 @@ fn write_node(file: &mut File, node: &Node) -> Result<(), Box<dyn std::error::Er
             NodeFieldType::OptionalLocation => {
                 writeln!(file, "    pub fn {}(&self) -> Option<Location<'pr>> {{", field.name)?;
                 writeln!(file, "        let pointer: *mut yp_location_t = unsafe {{ &mut (*self.pointer).{} }};", field.name)?;
-                writeln!(file, "        if pointer.is_null() {{")?;
+                writeln!(file, "        let start = unsafe {{ (*pointer).start }};")?;
+                writeln!(file, "        if start.is_null() {{")?;
                 writeln!(file, "            None")?;
                 writeln!(file, "        }} else {{")?;
                 writeln!(file, "            Some(Location {{ pointer: unsafe {{ NonNull::new_unchecked(pointer) }}, marker: PhantomData }})")?;
diff --git a/rust/yarp/src/lib.rs b/rust/yarp/src/lib.rs
index 25e06cc5661..9c5809b00e8 100644
--- a/rust/yarp/src/lib.rs
+++ b/rust/yarp/src/lib.rs
@@ -348,4 +348,40 @@ mod tests {
 
         assert_eq!(locals[0].as_slice(), b"x");
     }
+
+    #[test]
+    fn optional_loc_test() {
+        let source = r#"
+module Example
+  x = call_func(3, 4)
+  y = x.call_func 5, 6
+end
+"#;
+        let result = parse(source.as_ref());
+
+        let node = result.node();
+        let module = node.as_program_node().unwrap().statements().body().iter().next().unwrap();
+        let module = module.as_module_node().unwrap();
+        let body = module.body();
+        let writes = body.iter().next().unwrap().as_statements_node().unwrap().body().iter().collect::<Vec<_>>();
+        assert_eq!(writes.len(), 2);
+
+        let asgn = &writes[0];
+        let call = asgn.as_local_variable_write_node().unwrap().value();
+        let call = call.as_call_node().unwrap();
+
+        let operator_loc = call.operator_loc();
+        assert!(operator_loc.is_none());
+        let closing_loc = call.closing_loc();
+        assert!(closing_loc.is_some());
+
+        let asgn = &writes[1];
+        let call = asgn.as_local_variable_write_node().unwrap().value();
+        let call = call.as_call_node().unwrap();
+
+        let operator_loc = call.operator_loc();
+        assert!(operator_loc.is_some());
+        let closing_loc = call.closing_loc();
+        assert!(closing_loc.is_none());
+    }
 }