github.com/awesome-flow/flow@v0.0.3-0.20190918184116-508d75d68a2c/pkg/util/data/topology.go (about)

     1  package data
     2  
     3  import (
     4  	"fmt"
     5  )
     6  
     7  type TopologyNode interface{}
     8  
     9  type TopologyEdge struct {
    10  	From TopologyNode
    11  	To   TopologyNode
    12  }
    13  
    14  type Topology struct {
    15  	Edges map[TopologyEdge]struct{}
    16  	Nodes map[TopologyNode]struct{}
    17  }
    18  
    19  func NewTopology(nodes ...TopologyNode) *Topology {
    20  	nodeset := make(map[TopologyNode]struct{})
    21  	for _, node := range nodes {
    22  		nodeset[node] = struct{}{}
    23  	}
    24  	return &Topology{
    25  		Edges: make(map[TopologyEdge]struct{}),
    26  		Nodes: nodeset,
    27  	}
    28  }
    29  
    30  func (top *Topology) AddNode(node TopologyNode) {
    31  	top.Nodes[node] = struct{}{}
    32  }
    33  
    34  // ConnectTo creates a directed edge between node "from" to node "to".
    35  // If a topology represents a dependency graph, this notation should be
    36  // interpreted as: node "from" depends on node "to", i.e. in a case of a
    37  // topological sort node "to" would be visited before node "from".
    38  //
    39  // Example
    40  //
    41  // top := NewTopology(A, B)
    42  // top.Connect(A, B) // A -> B
    43  // top.Sort() // returns {B, A}: B has been visited, which satisfies A
    44  // dependencies.
    45  func (top *Topology) Connect(from, to TopologyNode) error {
    46  	if _, ok := top.Nodes[from]; !ok {
    47  		return fmt.Errorf("Can not connect from unknown node: %#v", from)
    48  	}
    49  	if _, ok := top.Nodes[to]; !ok {
    50  		return fmt.Errorf("Can not connect to unknown node: %#v", to)
    51  	}
    52  	top.Edges[TopologyEdge{From: from, To: to}] = struct{}{}
    53  
    54  	return nil
    55  }
    56  
    57  func (top *Topology) Sort() ([]TopologyNode, error) {
    58  	temp := make(map[TopologyNode]bool)
    59  	perm := make(map[TopologyNode]bool)
    60  	outs := make(map[TopologyNode][]TopologyNode)
    61  	for edge := range top.Edges {
    62  		if _, ok := outs[edge.From]; !ok {
    63  			outs[edge.From] = make([]TopologyNode, 0, 1)
    64  		}
    65  		outs[edge.From] = append(outs[edge.From], edge.To)
    66  	}
    67  
    68  	var visitAll func([]TopologyNode) ([]TopologyNode, error)
    69  	visitAll = func(nodes []TopologyNode) ([]TopologyNode, error) {
    70  		res := make([]TopologyNode, 0)
    71  		for _, node := range nodes {
    72  			if perm[node] {
    73  				continue
    74  			}
    75  			if temp[node] {
    76  				return nil, fmt.Errorf("Detected graph cycle on node %#v", node)
    77  			}
    78  			temp[node] = true
    79  			if subs, ok := outs[node]; ok {
    80  				subsorted, err := visitAll(subs)
    81  				if err != nil {
    82  					return nil, err
    83  				}
    84  				res = append(res, subsorted...)
    85  			}
    86  			perm[node] = true
    87  			res = append(res, node)
    88  		}
    89  		return res, nil
    90  	}
    91  	nodes := make([]TopologyNode, 0, len(top.Nodes))
    92  	for node := range top.Nodes {
    93  		nodes = append(nodes, node)
    94  	}
    95  
    96  	if res, err := visitAll(nodes); err != nil {
    97  		return nil, err
    98  	} else {
    99  		return res, nil
   100  	}
   101  }