redblack.go

Documentation: github.com/sassoftware/relic/lib/redblack

     1  //
     2  // Copyright (c) SAS Institute Inc.
     3  //
     4  // Licensed under the Apache License, Version 2.0 (the "License");
     5  // you may not use this file except in compliance with the License.
     6  // You may obtain a copy of the License at
     7  //
     8  //     http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  //
    16  
    17  // Simple, incomplete red-black tree implementation meant only to rebuild the
    18  // directory tree of a CDF file.
    19  package redblack
    20  
    21  func New(less LessFunc) *Tree {
    22  	return &Tree{Less: less}
    23  }
    24  
    25  func (t *Tree) Insert(item interface{}) {
    26  	node := &Node{Item: item, Less: t.Less}
    27  	t.Root = t.Root.insert(node)
    28  	t.Count++
    29  }
    30  
    31  func (t *Tree) Nodes() []*Node {
    32  	if t.Root == nil {
    33  		return nil
    34  	}
    35  	ret := make([]*Node, 0, t.Count)
    36  	stack := []*Node{t.Root}
    37  	for len(stack) > 0 {
    38  		i := len(stack) - 1
    39  		node := stack[i]
    40  		stack = stack[:i]
    41  		ret = append(ret, node)
    42  		if node.Children[0] != nil {
    43  			stack = append(stack, node.Children[0])
    44  		}
    45  		if node.Children[1] != nil {
    46  			stack = append(stack, node.Children[1])
    47  		}
    48  	}
    49  	return ret
    50  }
    51  
    52  type LessFunc func(i, j interface{}) bool
    53  
    54  type Tree struct {
    55  	Root  *Node
    56  	Less  LessFunc
    57  	Count uint
    58  }
    59  
    60  type Node struct {
    61  	Item     interface{}
    62  	Less     LessFunc
    63  	Red      bool
    64  	Children [2]*Node
    65  }
    66  
    67  func (n *Node) isRed() bool {
    68  	return n != nil && n.Red
    69  }
    70  
    71  func (n *Node) rotate(dir int) *Node {
    72  	a := n.Children[1-dir]
    73  	n.Children[1-dir] = a.Children[dir]
    74  	a.Children[dir] = n
    75  	n.Red = true
    76  	a.Red = false
    77  	return a
    78  }
    79  
    80  func (n *Node) insert(a *Node) *Node {
    81  	if n == nil {
    82  		return a
    83  	}
    84  	dir := 0
    85  	if n.Less(n.Item, a.Item) {
    86  		dir = 1
    87  	}
    88  	n.Children[dir] = n.Children[dir].insert(a)
    89  	if !n.Children[dir].isRed() {
    90  		return n
    91  	} else if n.Children[1-dir].isRed() {
    92  		n.Red = true
    93  		n.Children[0].Red = false
    94  		n.Children[1].Red = false
    95  		return n
    96  	} else if n.Children[dir].Children[dir].isRed() {
    97  		return n.rotate(1 - dir)
    98  	} else if n.Children[dir].Children[1-dir].isRed() {
    99  		n.Children[dir] = n.Children[dir].rotate(dir)
   100  		return n.rotate(1 - dir)
   101  	} else {
   102  		return n
   103  	}
   104  }
   105
View as plain text