github.com/pachyderm/pachyderm@v1.13.4/src/server/pkg/dag/dag.go (about)

     1  package dag
     2  
     3  // DAG represents a directected acyclic graph
     4  type DAG struct {
     5  	parents  map[string][]string
     6  	children map[string][]string
     7  	leaves   map[string]bool
     8  }
     9  
    10  // NewDAG creates a DAG and populates it with the given nodes.
    11  func NewDAG(nodes map[string][]string) *DAG {
    12  	result := &DAG{
    13  		parents:  make(map[string][]string),
    14  		children: make(map[string][]string),
    15  		leaves:   make(map[string]bool),
    16  	}
    17  	for id, parents := range nodes {
    18  		result.NewNode(id, parents)
    19  	}
    20  	return result
    21  }
    22  
    23  // NewNode adds a node to d.
    24  func (d *DAG) NewNode(id string, parents []string) {
    25  	d.parents[id] = parents
    26  	for _, parentID := range parents {
    27  		d.children[parentID] = append(d.children[parentID], id)
    28  		d.leaves[parentID] = false
    29  	}
    30  	if _, ok := d.leaves[id]; !ok {
    31  		d.leaves[id] = true
    32  	}
    33  }
    34  
    35  // Sorted returns all nodes in a topologically sorted order
    36  func (d *DAG) Sorted() []string {
    37  	seen := make(map[string]bool)
    38  	var result []string
    39  	for id := range d.parents {
    40  		result = append(result, dfs(id, d.parents, seen)...)
    41  	}
    42  	return result
    43  }
    44  
    45  // Leaves returns a slice containing all leaves in d.
    46  func (d *DAG) Leaves() []string {
    47  	var result []string
    48  	for id, isLeaf := range d.leaves {
    49  		// isLeaf might be false, explicit mark nodes as non leaves
    50  		if isLeaf {
    51  			result = append(result, id)
    52  		}
    53  	}
    54  	return result
    55  }
    56  
    57  // Ancestors returns a slice containing all ancestors of a node, 'id',
    58  // in d which are a descendant of at least one of the nodes in 'from'.
    59  func (d *DAG) Ancestors(id string, from []string) []string {
    60  	seen := make(map[string]bool)
    61  	for _, fromID := range from {
    62  		seen[fromID] = true
    63  	}
    64  	return dfs(id, d.parents, seen)
    65  }
    66  
    67  // Descendants returns a slice containing all descendants of a node, 'id',
    68  // in d which are an ancestor of at least one of the nodes in 'to'.
    69  func (d *DAG) Descendants(id string, to []string) []string {
    70  	seen := make(map[string]bool)
    71  	for _, toID := range to {
    72  		seen[toID] = true
    73  	}
    74  	return bfs(id, d.children, seen)
    75  }
    76  
    77  // Ghosts returns nodes that were referenced as parents but never created.
    78  func (d *DAG) Ghosts() []string {
    79  	var result []string
    80  	for id := range d.children {
    81  		if _, ok := d.parents[id]; !ok {
    82  			result = append(result, id)
    83  		}
    84  	}
    85  	return result
    86  }
    87  
    88  func dfs(id string, edges map[string][]string, seen map[string]bool) []string {
    89  	if seen[id] {
    90  		return nil
    91  	}
    92  
    93  	seen[id] = true
    94  
    95  	var result []string
    96  	for _, nID := range edges[id] {
    97  		result = append(result, dfs(nID, edges, seen)...)
    98  	}
    99  	return append(result, id)
   100  }
   101  
   102  func bfs(id string, edges map[string][]string, seen map[string]bool) []string {
   103  	var result []string
   104  	queue := []string{id}
   105  	for len(queue) != 0 {
   106  		result = append(result, queue[0])
   107  		queue = queue[1:]
   108  		for _, nID := range edges[result[len(result)-1]] {
   109  			if !seen[nID] {
   110  				seen[nID] = true
   111  				queue = append(queue, nID)
   112  			}
   113  		}
   114  	}
   115  	return result
   116  }