' Licensed to the .NET Foundation under one or more agreements.
' The .NET Foundation licenses this file to you under the MIT license.
' See the LICENSE file in the project root for more information.
 
Imports System.Collections.Immutable
Imports System.Threading
Imports Microsoft.CodeAnalysis.PooledObjects
Imports Microsoft.CodeAnalysis.VisualBasic.Symbols
Imports Microsoft.CodeAnalysis.VisualBasic.Symbols.Retargeting
 
Namespace Microsoft.CodeAnalysis.VisualBasic
 
    Partial Public Class VisualBasicCompilation
 
        Private _lazyUsedAssemblyReferences As ConcurrentSet(Of AssemblySymbol)
        Private _usedAssemblyReferencesFrozen As Boolean
 
        Public Overrides Function GetUsedAssemblyReferences(Optional cancellationToken As CancellationToken = Nothing) As ImmutableArray(Of MetadataReference)
            Dim usedAssemblies As ConcurrentSet(Of AssemblySymbol) = GetCompleteSetOfUsedAssemblies(cancellationToken)
 
            If usedAssemblies Is Nothing Then
                Return ImmutableArray(Of MetadataReference).Empty
            End If
 
            ' Use stable ordering for the result, matching the order in References.
            Dim builder = ArrayBuilder(Of MetadataReference).GetInstance(usedAssemblies.Count)
 
            For Each reference In References
                If reference.Properties.Kind = MetadataImageKind.Assembly Then
                    Dim symbol As Symbol = GetBoundReferenceManager().GetReferencedAssemblySymbol(reference)
 
                    If symbol IsNot Nothing AndAlso usedAssemblies.Contains(DirectCast(symbol, AssemblySymbol)) Then
                        builder.Add(reference)
                    End If
                End If
            Next
 
            Return builder.ToImmutableAndFree()
        End Function
 
        Private Function GetCompleteSetOfUsedAssemblies(cancellationToken As CancellationToken) As ConcurrentSet(Of AssemblySymbol)
 
            If Not _usedAssemblyReferencesFrozen AndAlso Not Volatile.Read(_usedAssemblyReferencesFrozen) Then
 
                Dim diagnostics = BindingDiagnosticBag.GetConcurrentInstance()
                RoslynDebug.Assert(diagnostics.AccumulatesDiagnostics)
 
                GetDiagnosticsWithoutFiltering(CompilationStage.Declare, includeEarlierStages:=True, diagnostics, cancellationToken)
 
                Dim seenErrors As Boolean = diagnostics.HasAnyErrors()
 
                If Not seenErrors Then
                    diagnostics.DiagnosticBag.Clear()
                    GetDiagnosticsForAllMethodBodies(hasDeclarationErrors:=False, diagnostics, doLowering:=True, cancellationToken)
                    seenErrors = diagnostics.HasAnyErrors()
 
                    If Not seenErrors Then
                        AddUsedAssemblies(diagnostics.DependenciesBag)
                    End If
                End If
 
                CompleteTheSetOfUsedAssemblies(seenErrors, cancellationToken)
 
                diagnostics.Free()
            End If
 
            Return _lazyUsedAssemblyReferences
        End Function
 
        Private Sub AddUsedAssembly(dependency As AssemblySymbol, stack As ArrayBuilder(Of AssemblySymbol))
            If AddUsedAssembly(dependency) Then
                stack.Push(dependency)
            End If
        End Sub
 
        Private Sub AddReferencedAssemblies(assembly As AssemblySymbol, includeMainModule As Boolean, stack As ArrayBuilder(Of AssemblySymbol))
            For i As Integer = If(includeMainModule, 0, 1) To assembly.Modules.Length - 1
                For Each dependency In assembly.Modules(i).ReferencedAssemblySymbols
                    AddUsedAssembly(dependency, stack)
                Next
            Next
        End Sub
 
        Private Sub CompleteTheSetOfUsedAssemblies(seenErrors As Boolean, cancellationToken As CancellationToken)
 
            If _usedAssemblyReferencesFrozen OrElse Volatile.Read(_usedAssemblyReferencesFrozen) Then
                Return
            End If
 
            If seenErrors Then
                ' Add all referenced assemblies
                For Each assembly As AssemblySymbol In SourceModule.ReferencedAssemblySymbols
                    AddUsedAssembly(assembly)
                Next
            Else
                ' Assume that all assemblies used by the added modules are also used
                For i As Integer = 1 To SourceAssembly.Modules.Length - 1
                    For Each dependency In SourceAssembly.Modules(i).ReferencedAssemblySymbols
                        AddUsedAssembly(dependency)
                    Next
                Next
 
                If _usedAssemblyReferencesFrozen OrElse Volatile.Read(_usedAssemblyReferencesFrozen) Then
                    Return
                End If
 
                ' Assume that all assemblies used by the used assemblies are also used
                ' This, for example, takes care of including facade assemblies that forward types around.
                If _lazyUsedAssemblyReferences IsNot Nothing Then
                    SyncLock _lazyUsedAssemblyReferences
                        If _usedAssemblyReferencesFrozen OrElse Volatile.Read(_usedAssemblyReferencesFrozen) Then
                            Return
                        End If
 
                        Dim stack = ArrayBuilder(Of AssemblySymbol).GetInstance(_lazyUsedAssemblyReferences.Count)
                        stack.AddRange(_lazyUsedAssemblyReferences)
 
                        While stack.Count <> 0
                            Dim current As AssemblySymbol = stack.Pop()
                            Dim usedAssemblies As ConcurrentSet(Of AssemblySymbol)
 
                            Dim sourceAssembly = TryCast(current, SourceAssemblySymbol)
 
                            If sourceAssembly IsNot Nothing Then
                                ' The set of assemblies used by the referenced compilation feels Like
                                ' a reasonable approximation to the set of assembly references that would
                                ' be emitted into the resulting binary for that compilation. An alternative
                                ' would be to attempt to emit and get the exact set of emitted references
                                ' in case of success. This might be too slow though.
                                usedAssemblies = sourceAssembly.DeclaringCompilation.GetCompleteSetOfUsedAssemblies(cancellationToken)
                                If usedAssemblies IsNot Nothing Then
                                    For Each dependency As AssemblySymbol In usedAssemblies
                                        Debug.Assert(Not dependency.IsLinked)
                                        AddUsedAssembly(dependency, stack)
                                    Next
                                End If
 
                                Continue While
                            End If
 
                            Dim retargetingAssembly = TryCast(current, RetargetingAssemblySymbol)
 
                            If retargetingAssembly IsNot Nothing Then
                                usedAssemblies = retargetingAssembly.UnderlyingAssembly.DeclaringCompilation.GetCompleteSetOfUsedAssemblies(cancellationToken)
                                If usedAssemblies IsNot Nothing Then
                                    For Each underlyingDependency As AssemblySymbol In retargetingAssembly.UnderlyingAssembly.SourceModule.ReferencedAssemblySymbols
                                        If Not underlyingDependency.IsLinked AndAlso usedAssemblies.Contains(underlyingDependency) Then
                                            Dim dependency As AssemblySymbol = Nothing
 
                                            If Not DirectCast(retargetingAssembly.Modules(0), RetargetingModuleSymbol).RetargetingDefinitions(underlyingDependency, dependency) Then
                                                Debug.Assert(retargetingAssembly.Modules(0).ReferencedAssemblySymbols.Contains(underlyingDependency))
                                                dependency = underlyingDependency
                                            End If
 
                                            AddUsedAssembly(dependency, stack)
                                        End If
                                    Next
                                End If
 
                                AddReferencedAssemblies(retargetingAssembly, includeMainModule:=False, stack)
 
                                Continue While
                            End If
 
                            AddReferencedAssemblies(current, includeMainModule:=True, stack)
                        End While
 
                        stack.Free()
                    End SyncLock
                End If
 
                If SourceAssembly.CorLibrary IsNot Nothing Then
                    ' Add core library
                    AddUsedAssembly(sourceAssembly.CorLibrary)
                End If
            End If
 
            _usedAssemblyReferencesFrozen = True
        End Sub
 
        Friend Sub AddUsedAssemblies(assemblies As ICollection(Of AssemblySymbol))
            If Not assemblies.IsNullOrEmpty() Then
                For Each candidate In assemblies
                    AddUsedAssembly(candidate)
                Next
            End If
        End Sub
 
        Friend Function AddUsedAssembly(assembly As AssemblySymbol) As Boolean
 
            If assembly Is Nothing OrElse assembly Is SourceAssembly OrElse assembly.IsMissing Then
                Return False
            End If
 
            If _lazyUsedAssemblyReferences Is Nothing Then
                Interlocked.CompareExchange(_lazyUsedAssemblyReferences, New ConcurrentSet(Of AssemblySymbol)(), Nothing)
            End If
 
#If DEBUG Then
            Dim wasFrozen As Boolean = _usedAssemblyReferencesFrozen
#End If
            Dim added As Boolean = _lazyUsedAssemblyReferences.Add(assembly)
 
#If DEBUG Then
            Debug.Assert(Not added OrElse Not wasFrozen)
#End If
            Return added
        End Function
 
    End Class
End Namespace