From aa78c778d7df022559fae083627a4e61d2b413ec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Laurent=20Th=C3=A9venoux?= <thevenoux@adacore.com>
Date: Mon, 8 Apr 2024 15:34:27 +0200
Subject: [PATCH] Fix p_body_part_for_decl for when body part is in separate
 unit

When the body part of a declaration is in a separate unit,
`p_body_part_for_decl` should check for the next part in the separate
body explicitely since no `__nextpart` symbol exists for it in the
environement.

This check was only done for declarations included in public parts,
that change extends it to private parts.
---
 ada/ast.py                                    | 31 +++++++++----------
 .../body_part_for_decl_separate/a.adb         | 13 ++++++++
 .../body_part_for_decl_separate/a~p.adb       |  5 +++
 .../body_part_for_decl_separate/test.out      | 11 +++++++
 .../body_part_for_decl_separate/test.yaml     |  2 ++
 5 files changed, 46 insertions(+), 16 deletions(-)
 create mode 100644 testsuite/tests/properties/body_part_for_decl_separate/a.adb
 create mode 100644 testsuite/tests/properties/body_part_for_decl_separate/a~p.adb
 create mode 100644 testsuite/tests/properties/body_part_for_decl_separate/test.out
 create mode 100644 testsuite/tests/properties/body_part_for_decl_separate/test.yaml

diff --git a/ada/ast.py b/ada/ast.py
index e6bea27cc..976a78893 100644
--- a/ada/ast.py
+++ b/ada/ast.py
@@ -10859,24 +10859,23 @@ def next_part_for_decl():
             Not(default_next_part.is_null),
             default_next_part,
 
-            # Self is declared in a private part
-            decl_scope.is_a(T.PrivatePart),
-            parent_decl.next_part_for_decl.then(
-                lambda np: Entity.get_body_in_env(np.children_env)
+            decl_scope.is_a(T.PrivatePart, T.PublicPart),
+            Let(
+                lambda other_part=If(
+                    decl_scope.is_a(T.PrivatePart),
+                    parent_decl.next_part_for_decl,
+                    parent_decl.decl_private_part
+                ):
+                # Search in other part
+                other_part.then(
+                    lambda op: Entity.get_body_in_env(op.children_env)
+                )
+                # If not found, search in body
+                ._or(parent_decl.body_part_for_decl.then(
+                    lambda np: Entity.get_body_in_env(np.children_env)
+                ))
             ),
 
-            # Self is declared in a public part
-            decl_scope.is_a(T.PublicPart),
-
-            # Search in private part
-            parent_decl.decl_private_part.then(
-                lambda dpp: Entity.get_body_in_env(dpp.children_env),
-            )
-            # If not found, search in body
-            ._or(parent_decl.body_part_for_decl.then(
-                lambda np: Entity.get_body_in_env(np.children_env)
-            )),
-
             # No declarative scope: Bail out!
             decl_scope.is_null, No(T.Body.entity),
 
diff --git a/testsuite/tests/properties/body_part_for_decl_separate/a.adb b/testsuite/tests/properties/body_part_for_decl_separate/a.adb
new file mode 100644
index 000000000..a962be657
--- /dev/null
+++ b/testsuite/tests/properties/body_part_for_decl_separate/a.adb
@@ -0,0 +1,13 @@
+procedure A is
+   package P is
+      procedure Bar;
+      --% node.p_body_part_for_decl()
+   private
+      procedure Foo;
+      --% node.p_body_part_for_decl()
+   end P;
+
+   package body P is separate;
+begin
+   null;
+end A;
diff --git a/testsuite/tests/properties/body_part_for_decl_separate/a~p.adb b/testsuite/tests/properties/body_part_for_decl_separate/a~p.adb
new file mode 100644
index 000000000..d94e44b15
--- /dev/null
+++ b/testsuite/tests/properties/body_part_for_decl_separate/a~p.adb
@@ -0,0 +1,5 @@
+separate(A)
+package body P is
+   procedure Bar is null;
+   procedure Foo is null;
+end P;
diff --git a/testsuite/tests/properties/body_part_for_decl_separate/test.out b/testsuite/tests/properties/body_part_for_decl_separate/test.out
new file mode 100644
index 000000000..e84453205
--- /dev/null
+++ b/testsuite/tests/properties/body_part_for_decl_separate/test.out
@@ -0,0 +1,11 @@
+Working on node <SubpDecl ["Bar"] a.adb:3:7-3:21>
+=================================================
+
+Eval 'node.p_body_part_for_decl()'
+Result: <NullSubpDecl ["Bar"] a~p.adb:3:4-3:26>
+
+Working on node <SubpDecl ["Foo"] a.adb:6:7-6:21>
+=================================================
+
+Eval 'node.p_body_part_for_decl()'
+Result: <NullSubpDecl ["Foo"] a~p.adb:4:4-4:26>
diff --git a/testsuite/tests/properties/body_part_for_decl_separate/test.yaml b/testsuite/tests/properties/body_part_for_decl_separate/test.yaml
new file mode 100644
index 000000000..8d3ef04cb
--- /dev/null
+++ b/testsuite/tests/properties/body_part_for_decl_separate/test.yaml
@@ -0,0 +1,2 @@
+driver: inline-playground
+input_sources: [a.adb]