Topological Sort Extension Methods

Option Strict On
Option Explicit On

Imports System.Runtime.CompilerServices

Module Module1

    <Extension()>
    Public Function TopologicalSort(Of TKey, TNode)(this As IEnumerable(Of TNode),
                                                    keySelector As Func(Of TNode, TKey),
                                                    sucessorsSelector As Func(Of TNode, IEnumerable(Of TKey))) As IEnumerable(Of TNode)

        Dim result As New Queue(Of TNode)
        Dim lookup As ILookup(Of TKey, TNode) = this.ToLookup(keySelector)
        Dim nodeInfos As New Dictionary(Of TKey, Tuple(Of Integer, List(Of TKey)))

        'Constroi o grafo com base nas chaves
        Dim successorNodeInfo As Tuple(Of Integer, List(Of TKey)) = Nothing '(nPredecessors, Sucessors)
        For Each predecessorKey As TKey In lookup.Select(Function(l) l.Key)
            If Not nodeInfos.ContainsKey(predecessorKey) Then nodeInfos(predecessorKey) = Tuple.Create(0, New List(Of TKey))

            For Each successorKey As TKey In sucessorsSelector(lookup(predecessorKey).FirstOrDefault())
                If Not nodeInfos.ContainsKey(successorKey) Then nodeInfos(successorKey) = Tuple.Create(0, New List(Of TKey))

                If Not nodeInfos(predecessorKey).Item2.Contains(successorKey) Then
                    successorNodeInfo = nodeInfos(predecessorKey)
                    nodeInfos(predecessorKey) = Tuple.Create(successorNodeInfo.Item1, successorNodeInfo.Item2.Union({successorKey}).ToList())

                    successorNodeInfo = nodeInfos(successorKey)
                    nodeInfos(successorKey) = Tuple.Create(successorNodeInfo.Item1 + 1, successorNodeInfo.Item2)

                End If

            Next

        Next

        Dim key As TKey = Nothing
        Dim predecessorNodeInfo As Tuple(Of Integer, List(Of TKey)) = Nothing  '(nPredecessors, Sucessors)
        Dim unvisited As New Queue(Of TKey)(nodeInfos.Where(Function(kvp) kvp.Value.Item1 = 0).Select(Function(kvp) kvp.Key))
        While (unvisited.Count <> 0)
            key = unvisited.Dequeue()

            For Each node As TNode In lookup(key)
                result.Enqueue(node)
            Next

            predecessorNodeInfo = nodeInfos(key)
            nodeInfos.Remove(key)

            For Each sucessorKey As TKey In predecessorNodeInfo.Item2
                successorNodeInfo = nodeInfos(sucessorKey)
                successorNodeInfo = Tuple.Create(successorNodeInfo.Item1 - 1, successorNodeInfo.Item2)
                nodeInfos(sucessorKey) = successorNodeInfo

                If successorNodeInfo.Item1 = 0 Then unvisited.Enqueue(sucessorKey)

            Next

            predecessorNodeInfo.Item2.Clear()

        End While

        If nodeInfos.Count > 0 Then Throw New NotImplementedException("Circular graph cannot be sorted using topological sort")

        Return result

    End Function

    <Extension()>
    Public Function TopologicalSortDepthFirst(Of TNode, TKey)(this As IEnumerable(Of TNode),
                                                              keySelector As Func(Of TNode, TKey),
                                                              predecessorSelector As Func(Of TNode, IEnumerable(Of TKey))) As IEnumerable(Of TNode)


        'Para cada nó construir um dicionário de Chave->Nós com essa chave.
        Dim lookup As ILookup(Of TKey, TNode) = this.ToLookup(keySelector)

        Dim result As New Queue(Of TNode)
        Dim ancestors As New Stack(Of TKey)
        Dim unvisited As New Stack(Of TKey)(lookup.Select(Function(g) g.Key))

        While unvisited.Count > 0
            Dim nodeKey As TKey = unvisited.Peek()
            Dim look As IEnumerable(Of TNode) = lookup(nodeKey)

            Dim predecessors As IEnumerable(Of TKey) = Nothing
            If look.Count > 0 Then
                predecessors = predecessorSelector(look.FirstOrDefault())
            Else
                predecessors = New List(Of TKey)()

            End If

            If predecessors.Count > 0 Then
                If ancestors.Count = 0 OrElse Not ancestors.Peek.Equals(nodeKey) Then
                    If ancestors.Contains(nodeKey) Then Throw New Exception("Circular graph cannot be sorted using topological sort")

                    ancestors.Push(nodeKey)
                    For Each predecessor In predecessors
                        unvisited.Push(predecessor)
                    Next

                    Continue While

                End If
                ancestors.Pop()

            End If

            unvisited.Pop()
            For Each node As TNode In look
                If Not result.Contains(node) Then result.Enqueue(node)
            Next


        End While

        Return result

    End Function

End Module