Permalink
Fetching contributors…
Cannot retrieve contributors at this time
212 lines (185 sloc) 3.95 KB
package gorgonia
import (
"fmt"
"sort"
"unsafe"
"github.com/xtgo/set"
)
// Nodes is a slice of nodes, but it also acts as a set of nodes by implementing the Sort interface
type Nodes []*Node
// implement sort.Interface
func (ns Nodes) Len() int { return len(ns) }
func (ns Nodes) Less(i, j int) bool {
return uintptr(unsafe.Pointer(ns[i])) < uintptr(unsafe.Pointer(ns[j]))
}
func (ns Nodes) Swap(i, j int) { ns[i], ns[j] = ns[j], ns[i] }
// uses xtgo/set stuff
// Set returns a uniquifies slice. It mutates the slice.
func (ns Nodes) Set() Nodes {
sort.Sort(ns)
size := set.Uniq(ns)
ns = ns[:size]
return ns
}
// Add adds to set
func (ns Nodes) Add(n *Node) Nodes {
for _, node := range ns {
if node == n {
return ns
}
}
ns = append(ns, n)
return ns
}
// Contains checks if the wanted node is in the set
func (ns Nodes) Contains(want *Node) bool {
for _, n := range ns {
if n == want {
return true
}
}
return false
}
// Format implements fmt.Formatter, which allows Nodes to be differently formatted depending on the verbs
func (ns Nodes) Format(s fmt.State, c rune) {
delimiter := ", "
if s.Flag(' ') {
delimiter = " "
}
if s.Flag('+') {
delimiter = ", \n"
}
switch c {
case 'd':
s.Write([]byte("["))
for i, n := range ns {
fmt.Fprintf(s, "%x", n.ID())
if i < len(ns)-1 {
fmt.Fprintf(s, "%s", delimiter)
}
}
s.Write([]byte("]"))
case 'v', 's':
s.Write([]byte("["))
for i, n := range ns {
if s.Flag('#') {
fmt.Fprintf(s, "%s :: %v", n.Name(), n.t)
} else {
fmt.Fprintf(s, "%s", n.Name())
}
if i < len(ns)-1 {
fmt.Fprintf(s, "%s", delimiter)
}
}
s.Write([]byte("]"))
case 'Y':
s.Write([]byte("["))
for i, n := range ns {
fmt.Fprintf(s, "%v", n.t)
if i < len(ns)-1 {
fmt.Fprintf(s, "%s", delimiter)
}
}
s.Write([]byte("]"))
case 'P':
s.Write([]byte("["))
for i, n := range ns {
fmt.Fprintf(s, "%p", n)
if i < len(ns)-1 {
fmt.Fprintf(s, "%s", delimiter)
}
}
s.Write([]byte("]"))
}
}
// Difference is ns - other. Bear in mind it is NOT commutative
func (ns Nodes) Difference(other Nodes) Nodes {
sort.Sort(ns)
sort.Sort(other)
s := append(ns, other...)
count := set.Diff(s, len(ns))
return s[:count]
}
// Intersect performs an intersection with other Nodes
func (ns Nodes) Intersect(other Nodes) Nodes {
sort.Sort(ns)
sort.Sort(other)
s := append(ns, other...)
count := set.Inter(s, len(ns))
return s[:count]
}
// AllSameGraph returns true if all the nodes in the slice belong to the same graph. Note that constants do not have to belong to the same graph.
func (ns Nodes) AllSameGraph() bool {
if len(ns) == 0 {
return false
}
var g *ExprGraph
for _, n := range ns {
if !n.isConstant() {
g = n.g
break
}
}
for _, n := range ns {
if n.g != g && !n.isConstant() {
return false
}
}
return true
}
// Equals returns true if two Nodes are the same
func (ns Nodes) Equals(other Nodes) bool {
if len(ns) != len(other) {
return false
}
for _, n := range ns {
if !other.Contains(n) {
return false
}
}
return true
}
func (ns Nodes) mapSet() NodeSet { return NewNodeSet(ns...) }
func (ns Nodes) index(n *Node) int {
for i, node := range ns {
if node == n {
return i
}
}
return -1
}
func (ns Nodes) reverse() {
l := len(ns)
for i := l/2 - 1; i >= 0; i-- {
o := l - 1 - i
ns[i], ns[o] = ns[o], ns[i]
}
}
func (ns Nodes) replace(what, with *Node) Nodes {
for i, n := range ns {
if n == what {
ns[i] = with
}
}
return ns
}
var removers = make(map[string]int)
func (ns Nodes) remove(what *Node) Nodes {
for i := ns.index(what); i != -1; i = ns.index(what) {
copy(ns[i:], ns[i+1:])
ns[len(ns)-1] = nil // to prevent any unwanted references so things can be GC'd away
ns = ns[:len(ns)-1]
}
return ns
}
func (ns Nodes) dimSizers() []DimSizer {
retVal := borrowDimSizers(len(ns))
for i, n := range ns {
if s, ok := n.op.(sizeOp); ok {
retVal[i] = s
} else {
retVal[i] = n.shape
}
}
return retVal
}