From 3c459e5e2fe0a529e8dbf04bfb045ae5a51680c8 Mon Sep 17 00:00:00 2001
From: Romain Beguet <beguet@adacore.com>
Date: Fri, 27 Oct 2023 17:21:18 +0200
Subject: [PATCH] Rework initial_env of contract pragmas.

This makes sure there is parent-children relation between the pragma node
and the type or subprogram it is associated with, to make it easier to
know whether we are inside such a contract.
---
 ada/ast.py | 39 +++++++++++++++++----------------------
 1 file changed, 17 insertions(+), 22 deletions(-)

diff --git a/ada/ast.py b/ada/ast.py
index 07c4b4f65..12ce3ffb8 100644
--- a/ada/ast.py
+++ b/ada/ast.py
@@ -10156,26 +10156,6 @@ def is_ghost_code():
             'Assert', 'Assert_And_Cut', 'Assume', 'Loop_Invariant'
         )
 
-    @langkit_property()
-    def xref_initial_env():
-        """
-        Contract pragmas such as ``Precondition`` have full visibility on their
-        associated subprogram's formals although they are not in an internal
-        scope. We handle that by overriding this ``xref_initial_env`` property
-        which will make sure that name resolution uses the env returned by this
-        property when constructing xref_equations.
-        """
-        return If(
-            Entity.id.name_symbol.any_of(
-                "Pre", "Post", "Pre'Class", "Post'Class",
-                "Precondition", "Postcondition",
-                "Precondition'Class", "Postcondition'Class",
-                "Test_Case", "Contract_Cases", "Predicate"
-            ),
-            Entity.associated_entities.at(0).children_env,
-            Entity.children_env
-        )
-
     @langkit_property()
     def xref_equation():
         return Cond(
@@ -10487,11 +10467,26 @@ def associated_entities():
     def initial_env():
         """
         Return the initial env name for a pragma clause. We use the
-        Standard package for top level use clauses.
+        Standard package for top level use clauses. For contract pragmas such
+        as ``Precondition`` or ``Predicate``, we use the env of the entity the
+        pragma is associated with in order to properly resolve references to
+        formals or to the type's ``SyntheticTypePredicateObjectDecl`` instance.
         """
-        return If(
+        return Cond(
             Self.parent.parent.is_a(CompilationUnit),
             named_env('Standard'),
+
+            Self.as_bare_entity.id.name_symbol.any_of(
+                "Pre", "Post", "Pre'Class", "Post'Class",
+                "Precondition", "Postcondition",
+                "Precondition'Class", "Postcondition'Class",
+                "Test_Case", "Contract_Cases", "Predicate"
+            ),
+            Self.as_bare_entity.associated_entities.at(0).then(
+                lambda ent: direct_env(ent.children_env),
+                default_val=current_env()
+            ),
+
             current_env()
         )