From c4b8cf16427c681e33ecc19cece8cd50a4d728fa Mon Sep 17 00:00:00 2001
From: Leigh McCulloch <351529+leighmcculloch@users.noreply.github.com>
Date: Fri, 6 Dec 2024 23:06:31 +1000
Subject: [PATCH] change the approach to using a qualified qself

---
 soroban-sdk-macros/src/syn_ext.rs | 51 +++++++++++++++++--------------
 1 file changed, 28 insertions(+), 23 deletions(-)

diff --git a/soroban-sdk-macros/src/syn_ext.rs b/soroban-sdk-macros/src/syn_ext.rs
index 64273c01..42bcb166 100644
--- a/soroban-sdk-macros/src/syn_ext.rs
+++ b/soroban-sdk-macros/src/syn_ext.rs
@@ -1,12 +1,11 @@
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote, ToTokens};
-use std::collections::HashMap;
 use syn::{
     parse::{Parse, ParseStream},
     punctuated::Punctuated,
     token::Comma,
     AngleBracketedGenericArguments, Attribute, GenericArgument, Path, PathArguments, PathSegment,
-    ReturnType, Token, TypePath,
+    QSelf, ReturnType, Token, TypePath,
 };
 use syn::{
     spanned::Spanned, token::And, Error, FnArg, Ident, ImplItem, ImplItemFn, ItemImpl, ItemTrait,
@@ -217,14 +216,10 @@ fn unpack_result(typ: &Type) -> Option<(Type, Type)> {
 fn flatten_associated_items_in_impl_fns(imp: &mut ItemImpl) {
     // TODO: Flatten associated consts used in functions.
     // Flatten associated types used in functions.
-    let associated_types = imp
-        .items
-        .iter()
-        .filter_map(|item| match item {
-            ImplItem::Type(i) => Some((i.ident.clone(), i.ty.clone())),
-            _ => None,
-        })
-        .collect::<HashMap<_, _>>();
+    let Some((_, trait_, _)) = &imp.trait_ else {
+        return;
+    };
+    let self_ty = &*imp.self_ty;
     let fn_input_types = imp
         .items
         .iter_mut()
@@ -237,19 +232,29 @@ fn flatten_associated_items_in_impl_fns(imp: &mut ItemImpl) {
         })
         .flatten();
     for t in fn_input_types {
-        if let Type::Path(TypePath { qself: None, path }) = t.as_mut() {
-            let segments = &path.segments;
-            if segments.len() == 2
-                && segments.first() == Some(&PathSegment::from(format_ident!("Self")))
-            {
-                if let Some(PathSegment {
-                    arguments: PathArguments::None,
-                    ident,
-                }) = segments.get(1)
-                {
-                    if let Some(resolved_ty) = associated_types.get(ident) {
-                        *t.as_mut() = resolved_ty.clone();
-                    }
+        let span = t.span();
+        if let Type::Path(TypePath {
+            qself: qself @ None,
+            path,
+        }) = t.as_mut()
+        {
+            let segments = path.segments.clone();
+            if segments.first() == Some(&PathSegment::from(format_ident!("Self"))) {
+                *qself = Some(QSelf {
+                    lt_token: Token![<](span),
+                    ty: Box::new(self_ty.clone()),
+                    // The index of the first path segment outside the <>.
+                    // For example, would be 2 for <Type as path::Trait>::AssociatedItem.
+                    position: trait_.segments.len(),
+                    as_token: Some(Token![as](span)),
+                    gt_token: Token![>](span),
+                });
+                *path = trait_.clone();
+                // Add the original path segments after the trait, skipping
+                // the first one which is the Self.
+                for segment in segments.into_iter().skip(1) {
+                    path.segments.push_punct(Token![::](span));
+                    path.segments.push_value(segment);
                 }
             }
         }