github.com/yinchengtsinghua/golang-Eos-dpos-Ethereum@v0.0.0-20190121132951-92cc4225ed8e/swarm/bmt/bmt.go (about)

     1  
     2  //此源码被清华学神尹成大魔王专业翻译分析并修改
     3  //尹成QQ77025077
     4  //尹成微信18510341407
     5  //尹成所在QQ群721929980
     6  //尹成邮箱 yinc13@mails.tsinghua.edu.cn
     7  //尹成毕业于清华大学,微软区块链领域全球最有价值专家
     8  //https://mvp.microsoft.com/zh-cn/PublicProfile/4033620
     9  //
    10  //
    11  //
    12  //
    13  //
    14  //
    15  //
    16  //
    17  //
    18  //
    19  //
    20  //
    21  //
    22  //
    23  //
    24  
    25  //
    26  package bmt
    27  
    28  import (
    29  	"fmt"
    30  	"hash"
    31  	"strings"
    32  	"sync"
    33  	"sync/atomic"
    34  )
    35  
    36  /*
    37  
    38  
    39  
    40  
    41  
    42  
    43  
    44  
    45  
    46  
    47  
    48  
    49  
    50  
    51  
    52  
    53  
    54    
    55  
    56    
    57  
    58    
    59   
    60   
    61   
    62   
    63  */
    64  
    65  
    66  const (
    67  //
    68  //
    69  	PoolSize = 8
    70  )
    71  
    72  //
    73  //
    74  type BaseHasherFunc func() hash.Hash
    75  
    76  //
    77  //
    78  //
    79  //
    80  //
    81  //
    82  //
    83  //
    84  //
    85  //
    86  type Hasher struct {
    87  pool *TreePool //
    88  bmt  *tree     //
    89  }
    90  
    91  //
    92  //
    93  func New(p *TreePool) *Hasher {
    94  	return &Hasher{
    95  		pool: p,
    96  	}
    97  }
    98  
    99  //
   100  //
   101  //
   102  type TreePool struct {
   103  	lock         sync.Mutex
   104  c            chan *tree     //
   105  hasher       BaseHasherFunc //
   106  SegmentSize  int            //
   107  SegmentCount int            //
   108  Capacity     int            //
   109  Depth        int            //
   110  Size         int            //
   111  count        int            //
   112  zerohashes   [][]byte       //
   113  }
   114  
   115  //
   116  //
   117  func NewTreePool(hasher BaseHasherFunc, segmentCount, capacity int) *TreePool {
   118  //
   119  	depth := calculateDepthFor(segmentCount)
   120  	segmentSize := hasher().Size()
   121  	zerohashes := make([][]byte, depth+1)
   122  	zeros := make([]byte, segmentSize)
   123  	zerohashes[0] = zeros
   124  	h := hasher()
   125  	for i := 1; i < depth+1; i++ {
   126  		zeros = doSum(h, nil, zeros, zeros)
   127  		zerohashes[i] = zeros
   128  	}
   129  	return &TreePool{
   130  		c:            make(chan *tree, capacity),
   131  		hasher:       hasher,
   132  		SegmentSize:  segmentSize,
   133  		SegmentCount: segmentCount,
   134  		Capacity:     capacity,
   135  		Size:         segmentCount * segmentSize,
   136  		Depth:        depth,
   137  		zerohashes:   zerohashes,
   138  	}
   139  }
   140  
   141  //
   142  func (p *TreePool) Drain(n int) {
   143  	p.lock.Lock()
   144  	defer p.lock.Unlock()
   145  	for len(p.c) > n {
   146  		<-p.c
   147  		p.count--
   148  	}
   149  }
   150  
   151  //
   152  //
   153  //
   154  func (p *TreePool) reserve() *tree {
   155  	p.lock.Lock()
   156  	defer p.lock.Unlock()
   157  	var t *tree
   158  	if p.count == p.Capacity {
   159  		return <-p.c
   160  	}
   161  	select {
   162  	case t = <-p.c:
   163  	default:
   164  		t = newTree(p.SegmentSize, p.Depth, p.hasher)
   165  		p.count++
   166  	}
   167  	return t
   168  }
   169  
   170  //
   171  //
   172  func (p *TreePool) release(t *tree) {
   173  p.c <- t //
   174  }
   175  
   176  //
   177  //
   178  //
   179  //
   180  type tree struct {
   181  leaves  []*node     //
   182  cursor  int         //
   183  offset  int         //
   184  section []byte      //
   185  result  chan []byte //
   186  span    []byte      //
   187  }
   188  
   189  //
   190  type node struct {
   191  isLeft      bool      //
   192  parent      *node     //
   193  state       int32     //
   194  left, right []byte    //
   195  hasher      hash.Hash //
   196  }
   197  
   198  //
   199  func newNode(index int, parent *node, hasher hash.Hash) *node {
   200  	return &node{
   201  		parent: parent,
   202  		isLeft: index%2 == 0,
   203  		hasher: hasher,
   204  	}
   205  }
   206  
   207  //
   208  func (t *tree) draw(hash []byte) string {
   209  	var left, right []string
   210  	var anc []*node
   211  	for i, n := range t.leaves {
   212  		left = append(left, fmt.Sprintf("%v", hashstr(n.left)))
   213  		if i%2 == 0 {
   214  			anc = append(anc, n.parent)
   215  		}
   216  		right = append(right, fmt.Sprintf("%v", hashstr(n.right)))
   217  	}
   218  	anc = t.leaves
   219  	var hashes [][]string
   220  	for l := 0; len(anc) > 0; l++ {
   221  		var nodes []*node
   222  		hash := []string{""}
   223  		for i, n := range anc {
   224  			hash = append(hash, fmt.Sprintf("%v|%v", hashstr(n.left), hashstr(n.right)))
   225  			if i%2 == 0 && n.parent != nil {
   226  				nodes = append(nodes, n.parent)
   227  			}
   228  		}
   229  		hash = append(hash, "")
   230  		hashes = append(hashes, hash)
   231  		anc = nodes
   232  	}
   233  	hashes = append(hashes, []string{"", fmt.Sprintf("%v", hashstr(hash)), ""})
   234  	total := 60
   235  	del := "                             "
   236  	var rows []string
   237  	for i := len(hashes) - 1; i >= 0; i-- {
   238  		var textlen int
   239  		hash := hashes[i]
   240  		for _, s := range hash {
   241  			textlen += len(s)
   242  		}
   243  		if total < textlen {
   244  			total = textlen + len(hash)
   245  		}
   246  		delsize := (total - textlen) / (len(hash) - 1)
   247  		if delsize > len(del) {
   248  			delsize = len(del)
   249  		}
   250  		row := fmt.Sprintf("%v: %v", len(hashes)-i-1, strings.Join(hash, del[:delsize]))
   251  		rows = append(rows, row)
   252  
   253  	}
   254  	rows = append(rows, strings.Join(left, "  "))
   255  	rows = append(rows, strings.Join(right, "  "))
   256  	return strings.Join(rows, "\n") + "\n"
   257  }
   258  
   259  //
   260  //
   261  func newTree(segmentSize, depth int, hashfunc func() hash.Hash) *tree {
   262  	n := newNode(0, nil, hashfunc())
   263  	prevlevel := []*node{n}
   264  //
   265  //
   266  	count := 2
   267  	for level := depth - 2; level >= 0; level-- {
   268  		nodes := make([]*node, count)
   269  		for i := 0; i < count; i++ {
   270  			parent := prevlevel[i/2]
   271  			var hasher hash.Hash
   272  			if level == 0 {
   273  				hasher = hashfunc()
   274  			}
   275  			nodes[i] = newNode(i, parent, hasher)
   276  		}
   277  		prevlevel = nodes
   278  		count *= 2
   279  	}
   280  //
   281  	return &tree{
   282  		leaves:  prevlevel,
   283  		result:  make(chan []byte),
   284  		section: make([]byte, 2*segmentSize),
   285  	}
   286  }
   287  
   288  //
   289  
   290  //
   291  func (h *Hasher) Size() int {
   292  	return h.pool.SegmentSize
   293  }
   294  
   295  //
   296  func (h *Hasher) BlockSize() int {
   297  	return 2 * h.pool.SegmentSize
   298  }
   299  
   300  //
   301  //
   302  //
   303  //
   304  //
   305  func (h *Hasher) Sum(b []byte) (s []byte) {
   306  	t := h.getTree()
   307  //
   308  	go h.writeSection(t.cursor, t.section, true, true)
   309  //
   310  	s = <-t.result
   311  	span := t.span
   312  //
   313  	h.releaseTree()
   314  //
   315  	if len(span) == 0 {
   316  		return append(b, s...)
   317  	}
   318  	return doSum(h.pool.hasher(), b, span, s)
   319  }
   320  
   321  //
   322  
   323  //
   324  //
   325  func (h *Hasher) Write(b []byte) (int, error) {
   326  	l := len(b)
   327  	if l == 0 || l > h.pool.Size {
   328  		return 0, nil
   329  	}
   330  	t := h.getTree()
   331  	secsize := 2 * h.pool.SegmentSize
   332  //
   333  	smax := secsize - t.offset
   334  //
   335  	if t.offset < secsize {
   336  //
   337  		copy(t.section[t.offset:], b)
   338  //
   339  //
   340  		if smax == 0 {
   341  			smax = secsize
   342  		}
   343  		if l <= smax {
   344  			t.offset += l
   345  			return l, nil
   346  		}
   347  	} else {
   348  //
   349  		if t.cursor == h.pool.SegmentCount*2 {
   350  			return 0, nil
   351  		}
   352  	}
   353  //
   354  	for smax < l {
   355  //
   356  		go h.writeSection(t.cursor, t.section, true, false)
   357  //
   358  		t.section = make([]byte, secsize)
   359  //
   360  		copy(t.section, b[smax:])
   361  //
   362  		t.cursor++
   363  //
   364  		smax += secsize
   365  	}
   366  	t.offset = l - smax + secsize
   367  	return l, nil
   368  }
   369  
   370  //
   371  func (h *Hasher) Reset() {
   372  	h.releaseTree()
   373  }
   374  
   375  //
   376  
   377  //
   378  //
   379  //
   380  func (h *Hasher) ResetWithLength(span []byte) {
   381  	h.Reset()
   382  	h.getTree().span = span
   383  }
   384  
   385  //
   386  //
   387  func (h *Hasher) releaseTree() {
   388  	t := h.bmt
   389  	if t == nil {
   390  		return
   391  	}
   392  	h.bmt = nil
   393  	go func() {
   394  		t.cursor = 0
   395  		t.offset = 0
   396  		t.span = nil
   397  		t.section = make([]byte, h.pool.SegmentSize*2)
   398  		select {
   399  		case <-t.result:
   400  		default:
   401  		}
   402  		h.pool.release(t)
   403  	}()
   404  }
   405  
   406  //
   407  func (h *Hasher) NewAsyncWriter(double bool) *AsyncHasher {
   408  	secsize := h.pool.SegmentSize
   409  	if double {
   410  		secsize *= 2
   411  	}
   412  	write := func(i int, section []byte, final bool) {
   413  		h.writeSection(i, section, double, final)
   414  	}
   415  	return &AsyncHasher{
   416  		Hasher:  h,
   417  		double:  double,
   418  		secsize: secsize,
   419  		write:   write,
   420  	}
   421  }
   422  
   423  //
   424  type SectionWriter interface {
   425  Reset()                                       //
   426  Write(index int, data []byte)                 //
   427  Sum(b []byte, length int, span []byte) []byte //
   428  SectionSize() int                             //
   429  }
   430  
   431  //
   432  //
   433  //
   434  //
   435  //
   436  //
   437  //
   438  //
   439  //
   440  //
   441  //
   442  //
   443  //
   444  //
   445  type AsyncHasher struct {
   446  *Hasher            //
   447  mtx     sync.Mutex //
   448  double  bool       //
   449  secsize int        //
   450  	write   func(i int, section []byte, final bool)
   451  }
   452  
   453  //
   454  
   455  //
   456  func (sw *AsyncHasher) SectionSize() int {
   457  	return sw.secsize
   458  }
   459  
   460  //
   461  //
   462  //
   463  func (sw *AsyncHasher) Write(i int, section []byte) {
   464  	sw.mtx.Lock()
   465  	defer sw.mtx.Unlock()
   466  	t := sw.getTree()
   467  //
   468  //
   469  	if i < t.cursor {
   470  //
   471  		go sw.write(i, section, false)
   472  		return
   473  	}
   474  //
   475  	if t.offset > 0 {
   476  		if i == t.cursor {
   477  //
   478  //
   479  			t.section = make([]byte, sw.secsize)
   480  			copy(t.section, section)
   481  			go sw.write(i, t.section, true)
   482  			return
   483  		}
   484  //
   485  		go sw.write(t.cursor, t.section, false)
   486  	}
   487  //
   488  //
   489  	t.cursor = i
   490  	t.offset = i*sw.secsize + 1
   491  	t.section = make([]byte, sw.secsize)
   492  	copy(t.section, section)
   493  }
   494  
   495  //
   496  //
   497  //
   498  //
   499  //
   500  //
   501  //
   502  //
   503  //
   504  func (sw *AsyncHasher) Sum(b []byte, length int, meta []byte) (s []byte) {
   505  	sw.mtx.Lock()
   506  	t := sw.getTree()
   507  	if length == 0 {
   508  		sw.mtx.Unlock()
   509  		s = sw.pool.zerohashes[sw.pool.Depth]
   510  	} else {
   511  //
   512  //
   513  		maxsec := (length - 1) / sw.secsize
   514  		if t.offset > 0 {
   515  			go sw.write(t.cursor, t.section, maxsec == t.cursor)
   516  		}
   517  //
   518  		t.cursor = maxsec
   519  		t.offset = length
   520  		result := t.result
   521  		sw.mtx.Unlock()
   522  //
   523  		s = <-result
   524  	}
   525  //
   526  	sw.releaseTree()
   527  //
   528  	if len(meta) == 0 {
   529  		return append(b, s...)
   530  	}
   531  //
   532  	return doSum(sw.pool.hasher(), b, meta, s)
   533  }
   534  
   535  //
   536  func (h *Hasher) writeSection(i int, section []byte, double bool, final bool) {
   537  //
   538  	var n *node
   539  	var isLeft bool
   540  	var hasher hash.Hash
   541  	var level int
   542  	t := h.getTree()
   543  	if double {
   544  		level++
   545  		n = t.leaves[i]
   546  		hasher = n.hasher
   547  		isLeft = n.isLeft
   548  		n = n.parent
   549  //
   550  		section = doSum(hasher, nil, section)
   551  	} else {
   552  		n = t.leaves[i/2]
   553  		hasher = n.hasher
   554  		isLeft = i%2 == 0
   555  	}
   556  //
   557  	if final {
   558  //
   559  		h.writeFinalNode(level, n, hasher, isLeft, section)
   560  	} else {
   561  		h.writeNode(n, hasher, isLeft, section)
   562  	}
   563  }
   564  
   565  //
   566  //
   567  //
   568  //
   569  //
   570  func (h *Hasher) writeNode(n *node, bh hash.Hash, isLeft bool, s []byte) {
   571  	level := 1
   572  	for {
   573  //
   574  		if n == nil {
   575  			h.getTree().result <- s
   576  			return
   577  		}
   578  //
   579  		if isLeft {
   580  			n.left = s
   581  		} else {
   582  			n.right = s
   583  		}
   584  //
   585  		if n.toggle() {
   586  			return
   587  		}
   588  //
   589  //
   590  		s = doSum(bh, nil, n.left, n.right)
   591  		isLeft = n.isLeft
   592  		n = n.parent
   593  		level++
   594  	}
   595  }
   596  
   597  //
   598  //
   599  //
   600  //
   601  //
   602  func (h *Hasher) writeFinalNode(level int, n *node, bh hash.Hash, isLeft bool, s []byte) {
   603  
   604  	for {
   605  //
   606  		if n == nil {
   607  			if s != nil {
   608  				h.getTree().result <- s
   609  			}
   610  			return
   611  		}
   612  		var noHash bool
   613  		if isLeft {
   614  //
   615  //
   616  //
   617  			n.right = h.pool.zerohashes[level]
   618  			if s != nil {
   619  				n.left = s
   620  //
   621  //
   622  //
   623  				noHash = false
   624  			} else {
   625  //
   626  				noHash = n.toggle()
   627  			}
   628  		} else {
   629  //
   630  			if s != nil {
   631  //
   632  				n.right = s
   633  //
   634  				noHash = n.toggle()
   635  
   636  			} else {
   637  //
   638  //
   639  				noHash = true
   640  			}
   641  		}
   642  //
   643  //
   644  //
   645  		if noHash {
   646  			s = nil
   647  		} else {
   648  			s = doSum(bh, nil, n.left, n.right)
   649  		}
   650  //
   651  		isLeft = n.isLeft
   652  		n = n.parent
   653  		level++
   654  	}
   655  }
   656  
   657  //
   658  func (h *Hasher) getTree() *tree {
   659  	if h.bmt != nil {
   660  		return h.bmt
   661  	}
   662  	t := h.pool.reserve()
   663  	h.bmt = t
   664  	return t
   665  }
   666  
   667  //
   668  //
   669  //
   670  func (n *node) toggle() bool {
   671  	return atomic.AddInt32(&n.state, 1)%2 == 1
   672  }
   673  
   674  //
   675  func doSum(h hash.Hash, b []byte, data ...[]byte) []byte {
   676  	h.Reset()
   677  	for _, v := range data {
   678  		h.Write(v)
   679  	}
   680  	return h.Sum(b)
   681  }
   682  
   683  //
   684  func hashstr(b []byte) string {
   685  	end := len(b)
   686  	if end > 4 {
   687  		end = 4
   688  	}
   689  	return fmt.Sprintf("%x", b[:end])
   690  }
   691  
   692  //
   693  func calculateDepthFor(n int) (d int) {
   694  	c := 2
   695  	for ; c < n; c *= 2 {
   696  		d++
   697  	}
   698  	return d + 1
   699  }