From 30f787ba8c2c8af03c47db8e7e26ecfe5bcf01e2 Mon Sep 17 00:00:00 2001
From: Romain Beguet <beguet@adacore.com>
Date: Wed, 26 Jun 2024 16:51:05 +0200
Subject: [PATCH] Make Annotate arguments resolution more robust.

---
 ada/ast.py                                    | 25 +++++++++++++------
 .../name_resolution/annotate_aspect/test.adb  |  5 ++++
 .../name_resolution/annotate_aspect/test.out  | 20 +++++++++++++++
 3 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/ada/ast.py b/ada/ast.py
index 9c8239374..6116df361 100644
--- a/ada/ast.py
+++ b/ada/ast.py
@@ -13760,15 +13760,24 @@ def annotate_argument_equation():
         Hence, we try here to resolve it using corresponding expected types.
         Note that we don't even check that the expected and actual types match
         in order to allow the same kind of flexibility that GNAT has.
+        GNAT also allows direct references to subprograms (even though they are
+        not valid expressions per-se), so we also fallback to resolving the
+        first visible declaration of the given name.
         """
-        return And(
-            Entity.sub_equation,
-            Self.expected_type_var.domain([
-                No(BaseTypeDecl.entity),
-                Self.std_string_type,
-                Self.std_wide_string_type,
-                Self.std_wide_wide_string_type
-            ])
+        return Or(
+            And(
+                Entity.sub_equation,
+                Self.expected_type_var.domain([
+                    No(BaseTypeDecl.entity),
+                    Self.std_string_type,
+                    Self.std_wide_string_type,
+                    Self.std_wide_wide_string_type
+                ])
+            ),
+            Entity.cast(Name).then(
+                lambda name: name.xref_no_overloading,
+                default_val=LogicFalse()
+            )
         )
 
     @langkit_property(public=True, dynamic_vars=[default_imprecise_fallback()],
diff --git a/testsuite/tests/name_resolution/annotate_aspect/test.adb b/testsuite/tests/name_resolution/annotate_aspect/test.adb
index d3d987514..cf6cbb1ef 100644
--- a/testsuite/tests/name_resolution/annotate_aspect/test.adb
+++ b/testsuite/tests/name_resolution/annotate_aspect/test.adb
@@ -21,6 +21,11 @@ procedure Test is
          "foo" & Wide_Wide_Character'Val (16#2665#),
          2 + 2);
    pragma Test_Block;
+
+   function Baz (X : Integer) return Integer is (X);
+   -- allow references to arbitrary declartations
+   pragma Annotate (My_Tool, My_Command, Baz);
+   pragma Test_Statement;
 begin
    null;
 end Test;
diff --git a/testsuite/tests/name_resolution/annotate_aspect/test.out b/testsuite/tests/name_resolution/annotate_aspect/test.out
index dde43bfed..0aa4063a2 100644
--- a/testsuite/tests/name_resolution/annotate_aspect/test.out
+++ b/testsuite/tests/name_resolution/annotate_aspect/test.out
@@ -270,5 +270,25 @@ Expr: <Int test.adb:22:14-22:15>
   type:          <ConcreteTypeDecl ["Universal_Int_Type_"] __standard:116:3-116:45>
   expected type: <ConcreteTypeDecl ["root_integer"] __standard:120:3-120:38>
 
+Resolving xrefs for node <PragmaNode test.adb:27:4-27:47>
+*********************************************************
+
+Expr: <Id "Annotate" test.adb:27:11-27:19>
+  references:    None
+  type:          None
+  expected type: None
+Expr: <Id "My_Tool" test.adb:27:21-27:28>
+  references:    None
+  type:          None
+  expected type: None
+Expr: <Id "My_Command" test.adb:27:30-27:40>
+  references:    None
+  type:          None
+  expected type: None
+Expr: <Id "Baz" test.adb:27:42-27:45>
+  references:    <DefiningName "Baz" test.adb:25:13-25:16>
+  type:          None
+  expected type: None
+
 
 Done.