77using System . Collections . Generic ;
88using System . Collections . Immutable ;
99using System . Diagnostics ;
10+ using System . Linq ;
1011using System . Text ;
1112using System . Threading ;
1213
@@ -24,7 +25,7 @@ public partial class Generator
2425 DiagnosticSeverity . Error ,
2526 isEnabledByDefault : true ) ;
2627
27- internal static readonly ImmutableArray < string > methods = ImmutableArray . Create ( new [ ]
28+ internal static readonly ImmutableHashSet < string > methods = ImmutableHashSet . Create ( new [ ]
2829 {
2930 // aggregation
3031 "Count" ,
@@ -98,20 +99,20 @@ public void Initialize(GeneratorInitializationContext context)
9899
99100 public void Execute ( GeneratorExecutionContext context )
100101 {
101- var typeSymbolCache = new TypeSymbolsCache ( context . Compilation ) ;
102+ var typeSymbolsCache = new TypeSymbolsCache ( context . Compilation ) ;
102103
103104 // Check if NetFabric.Hyperlinq.Abstractions and NetFabric.Hyperlinq.Abstractions are referenced
104- if ( typeSymbolCache [ "NetFabric.Hyperlinq.IValueEnumerable`2" ] is null
105- || typeSymbolCache [ "NetFabric.Hyperlinq.ValueEnumerableExtensions" ] is null )
106- return ;
105+ if ( typeSymbolsCache [ "NetFabric.Hyperlinq.IValueEnumerable`2" ] is null
106+ || typeSymbolsCache [ "NetFabric.Hyperlinq.ValueEnumerableExtensions" ] is null )
107+ return ; // TODO: return a Diagnostic?
107108
108109 if ( context . SyntaxReceiver is not SyntaxReceiver receiver )
109110 return ;
110111
111112 try
112113 {
113114 var builder = new CodeBuilder ( ) ;
114- GenerateSource ( context . Compilation , typeSymbolCache , receiver . MemberAccessExpressions , builder , context . CancellationToken ) ;
115+ GenerateSource ( context . Compilation , typeSymbolsCache , receiver . MemberAccessExpressions , builder , context . CancellationToken ) ;
115116 context . AddSource ( "ExtensionMethods.g.cs" , SourceText . From ( builder . ToString ( ) , Encoding . UTF8 ) ) ;
116117 }
117118 catch ( OperationCanceledException )
@@ -126,6 +127,8 @@ public void Execute(GeneratorExecutionContext context)
126127
127128 internal static void GenerateSource ( Compilation compilation , TypeSymbolsCache typeSymbolsCache , List < MemberAccessExpressionSyntax > memberAccessExpressions , CodeBuilder builder , CancellationToken cancellationToken , bool isUnitTest = false )
128129 {
130+ var generatedMethods = new HashSet < MethodSignature > ( ) ;
131+
129132 _ = builder
130133 . AppendLine ( "#nullable enable" )
131134 . AppendLine ( )
@@ -140,24 +143,26 @@ internal static void GenerateSource(Compilation compilation, TypeSymbolsCache ty
140143 {
141144 foreach ( var expressionSyntax in memberAccessExpressions )
142145 {
143- cancellationToken . ThrowIfCancellationRequested ( ) ;
146+ var semanticModel = compilation . GetSemanticModel ( expressionSyntax . SyntaxTree ) ;
144147
145- _ = expressionSyntax . Name . ToString ( ) switch
146- {
147- "AsValueEnumerable" => HandleAsValueEnumerable ( compilation , typeSymbolsCache , expressionSyntax , builder , cancellationToken , isUnitTest ) ,
148- _ => HandleMethod ( compilation , expressionSyntax , builder , cancellationToken , isUnitTest ) ,
149- } ;
148+ _ = GenerateSource ( compilation , semanticModel , typeSymbolsCache , expressionSyntax , builder , generatedMethods , cancellationToken , isUnitTest ) ;
150149 }
151150 }
152151 }
153152
154- static bool HandleMethod ( Compilation compilation , MemberAccessExpressionSyntax expressionSyntax , CodeBuilder builder , CancellationToken cancellationToken , bool isUnitTest )
153+ static ValueEnumerableType ? GenerateSource ( Compilation compilation , SemanticModel semanticModel , TypeSymbolsCache typeSymbolsCache , MemberAccessExpressionSyntax expressionSyntax , CodeBuilder builder , HashSet < MethodSignature > generatedMethods , CancellationToken cancellationToken , bool isUnitTest )
154+ => expressionSyntax . Name . ToString ( ) switch
155+ {
156+ "AsValueEnumerable" => GenerateAsValueEnumerable ( compilation , semanticModel , typeSymbolsCache , expressionSyntax , builder , generatedMethods , cancellationToken , isUnitTest ) ,
157+ _ => GenerateOperationSource ( compilation , semanticModel , expressionSyntax , builder , generatedMethods , cancellationToken , isUnitTest ) ,
158+ } ;
159+
160+ static ValueEnumerableType ? GenerateOperationSource ( Compilation compilation , SemanticModel semanticModel , MemberAccessExpressionSyntax expressionSyntax , CodeBuilder builder , HashSet < MethodSignature > generatedMethods , CancellationToken cancellationToken , bool isUnitTest )
155161 {
156- var semanticModel = compilation . GetSemanticModel ( expressionSyntax . SyntaxTree ) ;
157- if ( semanticModel . GetSymbolInfo ( expressionSyntax , cancellationToken ) . Symbol is not null ) // method already exists
158- return false ;
162+ // Get the type this operator is applied to
163+ var receiverTypeSymbol = semanticModel . GetTypeInfo ( expressionSyntax . Expression , cancellationToken ) . Type ;
159164
160- return false ;
165+ return null ;
161166 }
162167
163168 static bool IsValueEnumerable ( ITypeSymbol symbol , TypeSymbolsCache typeSymbolsCache )
0 commit comments