Avoid some capture variable wrapping in non-loops (#12235)

* avoid some capture variable wrapping in non-loops

* refine some more

* maybe this was important

src/filters/safe/capturedVars.ml

@@ -225,8 +225,26 @@ let captured_vars scom impl e =
 	and all_vars e =
 		let vars = ref PMap.empty in
 		let used = ref PMap.empty in
+		let in_loop = ref false in
+		let loop_vars = ref PMap.empty in
 		let assigned = ref PMap.empty in
 		let depth = ref 0 in
+		let mark_assigned v =
+			assigned := PMap.add v.v_id v !assigned;
+			(* Remove from vars lookup because we're done with it *)
+			vars := PMap.remove v.v_id !vars;
+		in
+		let check_loop_var v is_assign =
+			try
+				let is_assign' = PMap.find v.v_id !loop_vars in
+				(* If we have both read and write in the same loop we need to wrap. *)
+				if is_assign <> is_assign' then mark_assigned v
+			with Not_found ->
+				loop_vars := PMap.add v.v_id is_assign !loop_vars
+		in
+		let check_loop_var v is_assign =
+			if !in_loop then check_loop_var v is_assign
+		in
 		let rec collect_vars = function
 		| Block f ->
 			let old = !vars in
@@ -234,7 +252,12 @@ let captured_vars scom impl e =
 			vars := old;
 		| Loop f ->
 			let old = !vars in
+			let old_loop_vars = !loop_vars in
+			let old_loop = !in_loop in
+			in_loop := true;
 			f collect_vars;
+			in_loop := old_loop;
+			loop_vars := old_loop_vars;
 			vars := old;
 		| Function f ->
 			incr depth;
@@ -246,8 +269,8 @@ let captured_vars scom impl e =
 			(try
 				let d = PMap.find v.v_id !vars in
 				if d <> !depth then begin
+					check_loop_var v false;
 					used := PMap.add v.v_id v !used;
-					if has_var_flag v VAssigned then assigned := PMap.add v.v_id v !assigned;
 				end
 			with Not_found -> ())
 		| Assign v ->
@@ -256,13 +279,12 @@ let captured_vars scom impl e =
 				(* different depth - needs wrap *)
 				if d <> !depth then begin
 					used := PMap.add v.v_id v !used;
-					assigned := PMap.add v.v_id v !assigned;
-				end
+					mark_assigned v
 				(* same depth but assigned after being used on a different depth - needs wrap *)
-				else if PMap.mem v.v_id !used then
-					assigned := PMap.add v.v_id v !assigned
+				end else if PMap.mem v.v_id !used then
+					mark_assigned v
 				else
-					add_var_flag v VAssigned;
+					check_loop_var v true
 			with Not_found -> ())
 		in
 		local_usage collect_vars e;