tree_math_test.go raw

   1  package mls
   2  
   3  import (
   4  	"fmt"
   5  	"testing"
   6  )
   7  
   8  type treeMathTest struct {
   9  	NLeaves numLeaves `json:"n_leaves"`
  10  
  11  	NNodes  uint32       `json:"n_nodes"`
  12  	Root    nodeIndex    `json:"root"`
  13  	Left    []*nodeIndex `json:"left"`
  14  	Right   []*nodeIndex `json:"right"`
  15  	Parent  []*nodeIndex `json:"parent"`
  16  	Sibling []*nodeIndex `json:"sibling"`
  17  }
  18  
  19  func testTreeMath(t *testing.T, tc *treeMathTest) {
  20  	n := tc.NLeaves
  21  	if w := n.width(); w != tc.NNodes {
  22  		t.Errorf("width(%v) = %v, want %v", n, w, tc.NNodes)
  23  	}
  24  	if r := n.root(); r != tc.Root {
  25  		t.Errorf("root(%v) = %v, want %v", n, r, tc.Root)
  26  	}
  27  	for i, want := range tc.Left {
  28  		x := nodeIndex(i)
  29  		l := newOptionalNodeIndex(x.left())
  30  		if !optionalNodeIndexEqual(l, want) {
  31  			t.Errorf("left(%v) = %v, want %v", x, l, want)
  32  		}
  33  	}
  34  	for i, want := range tc.Right {
  35  		x := nodeIndex(i)
  36  		r := newOptionalNodeIndex(x.right())
  37  		if !optionalNodeIndexEqual(r, want) {
  38  			t.Errorf("right(%v) = %v, want %v", x, r, want)
  39  		}
  40  	}
  41  	for i, want := range tc.Parent {
  42  		x := nodeIndex(i)
  43  		p := newOptionalNodeIndex(n.parent(x))
  44  		if !optionalNodeIndexEqual(p, want) {
  45  			t.Errorf("parent(%v) = %v, want %v", x, p, want)
  46  		}
  47  	}
  48  	for i, want := range tc.Sibling {
  49  		x := nodeIndex(i)
  50  		s := newOptionalNodeIndex(n.sibling(x))
  51  		if !optionalNodeIndexEqual(s, want) {
  52  			t.Errorf("sibling(%v) = %v, want %v", x, s, want)
  53  		}
  54  	}
  55  }
  56  
  57  func newOptionalNodeIndex(x nodeIndex, ok bool) *nodeIndex {
  58  	if !ok {
  59  		return nil
  60  	}
  61  	return &x
  62  }
  63  
  64  func optionalNodeIndexEqual(x, y *nodeIndex) bool {
  65  	if x == nil || y == nil {
  66  		return x == nil && y == nil
  67  	}
  68  	return *x == *y
  69  }
  70  
  71  func TestTreeMath(t *testing.T) {
  72  	var tests []treeMathTest
  73  	loadTestVector(t, "testdata/tree-math.json", &tests)
  74  
  75  	for _, tc := range tests {
  76  		t.Run(fmt.Sprintf("numLeaves(%v)", tc.NLeaves), func(t *testing.T) {
  77  			testTreeMath(t, &tc)
  78  		})
  79  	}
  80  }
  81