cubefs

1
// Copyright 2018 Klaus Post. All rights reserved.
2
// Use of this source code is governed by a BSD-style
3
// license that can be found in the LICENSE file.
4
// Based on work Copyright (c) 2013, Yann Collet, released under BSD License.
5

6
package fse
7

8
import (
9
	"errors"
10
	"fmt"
11
)
12

13
// Compress the input bytes. Input must be < 2GB.
14
// Provide a Scratch buffer to avoid memory allocations.
15
// Note that the output is also kept in the scratch buffer.
16
// If input is too hard to compress, ErrIncompressible is returned.
17
// If input is a single byte value repeated ErrUseRLE is returned.
18
func Compress(in []byte, s *Scratch) ([]byte, error) {
19
	if len(in) <= 1 {
20
		return nil, ErrIncompressible
21
	}
22
	if len(in) > (2<<30)-1 {
23
		return nil, errors.New("input too big, must be < 2GB")
24
	}
25
	s, err := s.prepare(in)
26
	if err != nil {
27
		return nil, err
28
	}
29

30
	// Create histogram, if none was provided.
31
	maxCount := s.maxCount
32
	if maxCount == 0 {
33
		maxCount = s.countSimple(in)
34
	}
35
	// Reset for next run.
36
	s.clearCount = true
37
	s.maxCount = 0
38
	if maxCount == len(in) {
39
		// One symbol, use RLE
40
		return nil, ErrUseRLE
41
	}
42
	if maxCount == 1 || maxCount < (len(in)>>7) {
43
		// Each symbol present maximum once or too well distributed.
44
		return nil, ErrIncompressible
45
	}
46
	s.optimalTableLog()
47
	err = s.normalizeCount()
48
	if err != nil {
49
		return nil, err
50
	}
51
	err = s.writeCount()
52
	if err != nil {
53
		return nil, err
54
	}
55

56
	if false {
57
		err = s.validateNorm()
58
		if err != nil {
59
			return nil, err
60
		}
61
	}
62

63
	err = s.buildCTable()
64
	if err != nil {
65
		return nil, err
66
	}
67
	err = s.compress(in)
68
	if err != nil {
69
		return nil, err
70
	}
71
	s.Out = s.bw.out
72
	// Check if we compressed.
73
	if len(s.Out) >= len(in) {
74
		return nil, ErrIncompressible
75
	}
76
	return s.Out, nil
77
}
78

79
// cState contains the compression state of a stream.
80
type cState struct {
81
	bw         *bitWriter
82
	stateTable []uint16
83
	state      uint16
84
}
85

86
// init will initialize the compression state to the first symbol of the stream.
87
func (c *cState) init(bw *bitWriter, ct *cTable, tableLog uint8, first symbolTransform) {
88
	c.bw = bw
89
	c.stateTable = ct.stateTable
90

91
	nbBitsOut := (first.deltaNbBits + (1 << 15)) >> 16
92
	im := int32((nbBitsOut << 16) - first.deltaNbBits)
93
	lu := (im >> nbBitsOut) + first.deltaFindState
94
	c.state = c.stateTable[lu]
95
}
96

97
// encode the output symbol provided and write it to the bitstream.
98
func (c *cState) encode(symbolTT symbolTransform) {
99
	nbBitsOut := (uint32(c.state) + symbolTT.deltaNbBits) >> 16
100
	dstState := int32(c.state>>(nbBitsOut&15)) + symbolTT.deltaFindState
101
	c.bw.addBits16NC(c.state, uint8(nbBitsOut))
102
	c.state = c.stateTable[dstState]
103
}
104

105
// encode the output symbol provided and write it to the bitstream.
106
func (c *cState) encodeZero(symbolTT symbolTransform) {
107
	nbBitsOut := (uint32(c.state) + symbolTT.deltaNbBits) >> 16
108
	dstState := int32(c.state>>(nbBitsOut&15)) + symbolTT.deltaFindState
109
	c.bw.addBits16ZeroNC(c.state, uint8(nbBitsOut))
110
	c.state = c.stateTable[dstState]
111
}
112

113
// flush will write the tablelog to the output and flush the remaining full bytes.
114
func (c *cState) flush(tableLog uint8) {
115
	c.bw.flush32()
116
	c.bw.addBits16NC(c.state, tableLog)
117
	c.bw.flush()
118
}
119

120
// compress is the main compression loop that will encode the input from the last byte to the first.
121
func (s *Scratch) compress(src []byte) error {
122
	if len(src) <= 2 {
123
		return errors.New("compress: src too small")
124
	}
125
	tt := s.ct.symbolTT[:256]
126
	s.bw.reset(s.Out)
127

128
	// Our two states each encodes every second byte.
129
	// Last byte encoded (first byte decoded) will always be encoded by c1.
130
	var c1, c2 cState
131

132
	// Encode so remaining size is divisible by 4.
133
	ip := len(src)
134
	if ip&1 == 1 {
135
		c1.init(&s.bw, &s.ct, s.actualTableLog, tt[src[ip-1]])
136
		c2.init(&s.bw, &s.ct, s.actualTableLog, tt[src[ip-2]])
137
		c1.encodeZero(tt[src[ip-3]])
138
		ip -= 3
139
	} else {
140
		c2.init(&s.bw, &s.ct, s.actualTableLog, tt[src[ip-1]])
141
		c1.init(&s.bw, &s.ct, s.actualTableLog, tt[src[ip-2]])
142
		ip -= 2
143
	}
144
	if ip&2 != 0 {
145
		c2.encodeZero(tt[src[ip-1]])
146
		c1.encodeZero(tt[src[ip-2]])
147
		ip -= 2
148
	}
149

150
	// Main compression loop.
151
	switch {
152
	case !s.zeroBits && s.actualTableLog <= 8:
153
		// We can encode 4 symbols without requiring a flush.
154
		// We do not need to check if any output is 0 bits.
155
		for ip >= 4 {
156
			s.bw.flush32()
157
			v3, v2, v1, v0 := src[ip-4], src[ip-3], src[ip-2], src[ip-1]
158
			c2.encode(tt[v0])
159
			c1.encode(tt[v1])
160
			c2.encode(tt[v2])
161
			c1.encode(tt[v3])
162
			ip -= 4
163
		}
164
	case !s.zeroBits:
165
		// We do not need to check if any output is 0 bits.
166
		for ip >= 4 {
167
			s.bw.flush32()
168
			v3, v2, v1, v0 := src[ip-4], src[ip-3], src[ip-2], src[ip-1]
169
			c2.encode(tt[v0])
170
			c1.encode(tt[v1])
171
			s.bw.flush32()
172
			c2.encode(tt[v2])
173
			c1.encode(tt[v3])
174
			ip -= 4
175
		}
176
	case s.actualTableLog <= 8:
177
		// We can encode 4 symbols without requiring a flush
178
		for ip >= 4 {
179
			s.bw.flush32()
180
			v3, v2, v1, v0 := src[ip-4], src[ip-3], src[ip-2], src[ip-1]
181
			c2.encodeZero(tt[v0])
182
			c1.encodeZero(tt[v1])
183
			c2.encodeZero(tt[v2])
184
			c1.encodeZero(tt[v3])
185
			ip -= 4
186
		}
187
	default:
188
		for ip >= 4 {
189
			s.bw.flush32()
190
			v3, v2, v1, v0 := src[ip-4], src[ip-3], src[ip-2], src[ip-1]
191
			c2.encodeZero(tt[v0])
192
			c1.encodeZero(tt[v1])
193
			s.bw.flush32()
194
			c2.encodeZero(tt[v2])
195
			c1.encodeZero(tt[v3])
196
			ip -= 4
197
		}
198
	}
199

200
	// Flush final state.
201
	// Used to initialize state when decoding.
202
	c2.flush(s.actualTableLog)
203
	c1.flush(s.actualTableLog)
204

205
	return s.bw.close()
206
}
207

208
// writeCount will write the normalized histogram count to header.
209
// This is read back by readNCount.
210
func (s *Scratch) writeCount() error {
211
	var (
212
		tableLog  = s.actualTableLog
213
		tableSize = 1 << tableLog
214
		previous0 bool
215
		charnum   uint16
216

217
		maxHeaderSize = ((int(s.symbolLen) * int(tableLog)) >> 3) + 3
218

219
		// Write Table Size
220
		bitStream = uint32(tableLog - minTablelog)
221
		bitCount  = uint(4)
222
		remaining = int16(tableSize + 1) /* +1 for extra accuracy */
223
		threshold = int16(tableSize)
224
		nbBits    = uint(tableLog + 1)
225
	)
226
	if cap(s.Out) < maxHeaderSize {
227
		s.Out = make([]byte, 0, s.br.remain()+maxHeaderSize)
228
	}
229
	outP := uint(0)
230
	out := s.Out[:maxHeaderSize]
231

232
	// stops at 1
233
	for remaining > 1 {
234
		if previous0 {
235
			start := charnum
236
			for s.norm[charnum] == 0 {
237
				charnum++
238
			}
239
			for charnum >= start+24 {
240
				start += 24
241
				bitStream += uint32(0xFFFF) << bitCount
242
				out[outP] = byte(bitStream)
243
				out[outP+1] = byte(bitStream >> 8)
244
				outP += 2
245
				bitStream >>= 16
246
			}
247
			for charnum >= start+3 {
248
				start += 3
249
				bitStream += 3 << bitCount
250
				bitCount += 2
251
			}
252
			bitStream += uint32(charnum-start) << bitCount
253
			bitCount += 2
254
			if bitCount > 16 {
255
				out[outP] = byte(bitStream)
256
				out[outP+1] = byte(bitStream >> 8)
257
				outP += 2
258
				bitStream >>= 16
259
				bitCount -= 16
260
			}
261
		}
262

263
		count := s.norm[charnum]
264
		charnum++
265
		max := (2*threshold - 1) - remaining
266
		if count < 0 {
267
			remaining += count
268
		} else {
269
			remaining -= count
270
		}
271
		count++ // +1 for extra accuracy
272
		if count >= threshold {
273
			count += max // [0..max[ [max..threshold[ (...) [threshold+max 2*threshold[
274
		}
275
		bitStream += uint32(count) << bitCount
276
		bitCount += nbBits
277
		if count < max {
278
			bitCount--
279
		}
280

281
		previous0 = count == 1
282
		if remaining < 1 {
283
			return errors.New("internal error: remaining<1")
284
		}
285
		for remaining < threshold {
286
			nbBits--
287
			threshold >>= 1
288
		}
289

290
		if bitCount > 16 {
291
			out[outP] = byte(bitStream)
292
			out[outP+1] = byte(bitStream >> 8)
293
			outP += 2
294
			bitStream >>= 16
295
			bitCount -= 16
296
		}
297
	}
298

299
	out[outP] = byte(bitStream)
300
	out[outP+1] = byte(bitStream >> 8)
301
	outP += (bitCount + 7) / 8
302

303
	if charnum > s.symbolLen {
304
		return errors.New("internal error: charnum > s.symbolLen")
305
	}
306
	s.Out = out[:outP]
307
	return nil
308
}
309

310
// symbolTransform contains the state transform for a symbol.
311
type symbolTransform struct {
312
	deltaFindState int32
313
	deltaNbBits    uint32
314
}
315

316
// String prints values as a human readable string.
317
func (s symbolTransform) String() string {
318
	return fmt.Sprintf("dnbits: %08x, fs:%d", s.deltaNbBits, s.deltaFindState)
319
}
320

321
// cTable contains tables used for compression.
322
type cTable struct {
323
	tableSymbol []byte
324
	stateTable  []uint16
325
	symbolTT    []symbolTransform
326
}
327

328
// allocCtable will allocate tables needed for compression.
329
// If existing tables a re big enough, they are simply re-used.
330
func (s *Scratch) allocCtable() {
331
	tableSize := 1 << s.actualTableLog
332
	// get tableSymbol that is big enough.
333
	if cap(s.ct.tableSymbol) < tableSize {
334
		s.ct.tableSymbol = make([]byte, tableSize)
335
	}
336
	s.ct.tableSymbol = s.ct.tableSymbol[:tableSize]
337

338
	ctSize := tableSize
339
	if cap(s.ct.stateTable) < ctSize {
340
		s.ct.stateTable = make([]uint16, ctSize)
341
	}
342
	s.ct.stateTable = s.ct.stateTable[:ctSize]
343

344
	if cap(s.ct.symbolTT) < 256 {
345
		s.ct.symbolTT = make([]symbolTransform, 256)
346
	}
347
	s.ct.symbolTT = s.ct.symbolTT[:256]
348
}
349

350
// buildCTable will populate the compression table so it is ready to be used.
351
func (s *Scratch) buildCTable() error {
352
	tableSize := uint32(1 << s.actualTableLog)
353
	highThreshold := tableSize - 1
354
	var cumul [maxSymbolValue + 2]int16
355

356
	s.allocCtable()
357
	tableSymbol := s.ct.tableSymbol[:tableSize]
358
	// symbol start positions
359
	{
360
		cumul[0] = 0
361
		for ui, v := range s.norm[:s.symbolLen-1] {
362
			u := byte(ui) // one less than reference
363
			if v == -1 {
364
				// Low proba symbol
365
				cumul[u+1] = cumul[u] + 1
366
				tableSymbol[highThreshold] = u
367
				highThreshold--
368
			} else {
369
				cumul[u+1] = cumul[u] + v
370
			}
371
		}
372
		// Encode last symbol separately to avoid overflowing u
373
		u := int(s.symbolLen - 1)
374
		v := s.norm[s.symbolLen-1]
375
		if v == -1 {
376
			// Low proba symbol
377
			cumul[u+1] = cumul[u] + 1
378
			tableSymbol[highThreshold] = byte(u)
379
			highThreshold--
380
		} else {
381
			cumul[u+1] = cumul[u] + v
382
		}
383
		if uint32(cumul[s.symbolLen]) != tableSize {
384
			return fmt.Errorf("internal error: expected cumul[s.symbolLen] (%d) == tableSize (%d)", cumul[s.symbolLen], tableSize)
385
		}
386
		cumul[s.symbolLen] = int16(tableSize) + 1
387
	}
388
	// Spread symbols
389
	s.zeroBits = false
390
	{
391
		step := tableStep(tableSize)
392
		tableMask := tableSize - 1
393
		var position uint32
394
		// if any symbol > largeLimit, we may have 0 bits output.
395
		largeLimit := int16(1 << (s.actualTableLog - 1))
396
		for ui, v := range s.norm[:s.symbolLen] {
397
			symbol := byte(ui)
398
			if v > largeLimit {
399
				s.zeroBits = true
400
			}
401
			for nbOccurrences := int16(0); nbOccurrences < v; nbOccurrences++ {
402
				tableSymbol[position] = symbol
403
				position = (position + step) & tableMask
404
				for position > highThreshold {
405
					position = (position + step) & tableMask
406
				} /* Low proba area */
407
			}
408
		}
409

410
		// Check if we have gone through all positions
411
		if position != 0 {
412
			return errors.New("position!=0")
413
		}
414
	}
415

416
	// Build table
417
	table := s.ct.stateTable
418
	{
419
		tsi := int(tableSize)
420
		for u, v := range tableSymbol {
421
			// TableU16 : sorted by symbol order; gives next state value
422
			table[cumul[v]] = uint16(tsi + u)
423
			cumul[v]++
424
		}
425
	}
426

427
	// Build Symbol Transformation Table
428
	{
429
		total := int16(0)
430
		symbolTT := s.ct.symbolTT[:s.symbolLen]
431
		tableLog := s.actualTableLog
432
		tl := (uint32(tableLog) << 16) - (1 << tableLog)
433
		for i, v := range s.norm[:s.symbolLen] {
434
			switch v {
435
			case 0:
436
			case -1, 1:
437
				symbolTT[i].deltaNbBits = tl
438
				symbolTT[i].deltaFindState = int32(total - 1)
439
				total++
440
			default:
441
				maxBitsOut := uint32(tableLog) - highBits(uint32(v-1))
442
				minStatePlus := uint32(v) << maxBitsOut
443
				symbolTT[i].deltaNbBits = (maxBitsOut << 16) - minStatePlus
444
				symbolTT[i].deltaFindState = int32(total - v)
445
				total += v
446
			}
447
		}
448
		if total != int16(tableSize) {
449
			return fmt.Errorf("total mismatch %d (got) != %d (want)", total, tableSize)
450
		}
451
	}
452
	return nil
453
}
454

455
// countSimple will create a simple histogram in s.count.
456
// Returns the biggest count.
457
// Does not update s.clearCount.
458
func (s *Scratch) countSimple(in []byte) (max int) {
459
	for _, v := range in {
460
		s.count[v]++
461
	}
462
	m := uint32(0)
463
	for i, v := range s.count[:] {
464
		if v > m {
465
			m = v
466
		}
467
		if v > 0 {
468
			s.symbolLen = uint16(i) + 1
469
		}
470
	}
471
	return int(m)
472
}
473

474
// minTableLog provides the minimum logSize to safely represent a distribution.
475
func (s *Scratch) minTableLog() uint8 {
476
	minBitsSrc := highBits(uint32(s.br.remain()-1)) + 1
477
	minBitsSymbols := highBits(uint32(s.symbolLen-1)) + 2
478
	if minBitsSrc < minBitsSymbols {
479
		return uint8(minBitsSrc)
480
	}
481
	return uint8(minBitsSymbols)
482
}
483

484
// optimalTableLog calculates and sets the optimal tableLog in s.actualTableLog
485
func (s *Scratch) optimalTableLog() {
486
	tableLog := s.TableLog
487
	minBits := s.minTableLog()
488
	maxBitsSrc := uint8(highBits(uint32(s.br.remain()-1))) - 2
489
	if maxBitsSrc < tableLog {
490
		// Accuracy can be reduced
491
		tableLog = maxBitsSrc
492
	}
493
	if minBits > tableLog {
494
		tableLog = minBits
495
	}
496
	// Need a minimum to safely represent all symbol values
497
	if tableLog < minTablelog {
498
		tableLog = minTablelog
499
	}
500
	if tableLog > maxTableLog {
501
		tableLog = maxTableLog
502
	}
503
	s.actualTableLog = tableLog
504
}
505

506
var rtbTable = [...]uint32{0, 473195, 504333, 520860, 550000, 700000, 750000, 830000}
507

508
// normalizeCount will normalize the count of the symbols so
509
// the total is equal to the table size.
510
func (s *Scratch) normalizeCount() error {
511
	var (
512
		tableLog          = s.actualTableLog
513
		scale             = 62 - uint64(tableLog)
514
		step              = (1 << 62) / uint64(s.br.remain())
515
		vStep             = uint64(1) << (scale - 20)
516
		stillToDistribute = int16(1 << tableLog)
517
		largest           int
518
		largestP          int16
519
		lowThreshold      = (uint32)(s.br.remain() >> tableLog)
520
	)
521

522
	for i, cnt := range s.count[:s.symbolLen] {
523
		// already handled
524
		// if (count[s] == s.length) return 0;   /* rle special case */
525

526
		if cnt == 0 {
527
			s.norm[i] = 0
528
			continue
529
		}
530
		if cnt <= lowThreshold {
531
			s.norm[i] = -1
532
			stillToDistribute--
533
		} else {
534
			proba := (int16)((uint64(cnt) * step) >> scale)
535
			if proba < 8 {
536
				restToBeat := vStep * uint64(rtbTable[proba])
537
				v := uint64(cnt)*step - (uint64(proba) << scale)
538
				if v > restToBeat {
539
					proba++
540
				}
541
			}
542
			if proba > largestP {
543
				largestP = proba
544
				largest = i
545
			}
546
			s.norm[i] = proba
547
			stillToDistribute -= proba
548
		}
549
	}
550

551
	if -stillToDistribute >= (s.norm[largest] >> 1) {
552
		// corner case, need another normalization method
553
		return s.normalizeCount2()
554
	}
555
	s.norm[largest] += stillToDistribute
556
	return nil
557
}
558

559
// Secondary normalization method.
560
// To be used when primary method fails.
561
func (s *Scratch) normalizeCount2() error {
562
	const notYetAssigned = -2
563
	var (
564
		distributed  uint32
565
		total        = uint32(s.br.remain())
566
		tableLog     = s.actualTableLog
567
		lowThreshold = total >> tableLog
568
		lowOne       = (total * 3) >> (tableLog + 1)
569
	)
570
	for i, cnt := range s.count[:s.symbolLen] {
571
		if cnt == 0 {
572
			s.norm[i] = 0
573
			continue
574
		}
575
		if cnt <= lowThreshold {
576
			s.norm[i] = -1
577
			distributed++
578
			total -= cnt
579
			continue
580
		}
581
		if cnt <= lowOne {
582
			s.norm[i] = 1
583
			distributed++
584
			total -= cnt
585
			continue
586
		}
587
		s.norm[i] = notYetAssigned
588
	}
589
	toDistribute := (1 << tableLog) - distributed
590

591
	if (total / toDistribute) > lowOne {
592
		// risk of rounding to zero
593
		lowOne = (total * 3) / (toDistribute * 2)
594
		for i, cnt := range s.count[:s.symbolLen] {
595
			if (s.norm[i] == notYetAssigned) && (cnt <= lowOne) {
596
				s.norm[i] = 1
597
				distributed++
598
				total -= cnt
599
				continue
600
			}
601
		}
602
		toDistribute = (1 << tableLog) - distributed
603
	}
604
	if distributed == uint32(s.symbolLen)+1 {
605
		// all values are pretty poor;
606
		//   probably incompressible data (should have already been detected);
607
		//   find max, then give all remaining points to max
608
		var maxV int
609
		var maxC uint32
610
		for i, cnt := range s.count[:s.symbolLen] {
611
			if cnt > maxC {
612
				maxV = i
613
				maxC = cnt
614
			}
615
		}
616
		s.norm[maxV] += int16(toDistribute)
617
		return nil
618
	}
619

620
	if total == 0 {
621
		// all of the symbols were low enough for the lowOne or lowThreshold
622
		for i := uint32(0); toDistribute > 0; i = (i + 1) % (uint32(s.symbolLen)) {
623
			if s.norm[i] > 0 {
624
				toDistribute--
625
				s.norm[i]++
626
			}
627
		}
628
		return nil
629
	}
630

631
	var (
632
		vStepLog = 62 - uint64(tableLog)
633
		mid      = uint64((1 << (vStepLog - 1)) - 1)
634
		rStep    = (((1 << vStepLog) * uint64(toDistribute)) + mid) / uint64(total) // scale on remaining
635
		tmpTotal = mid
636
	)
637
	for i, cnt := range s.count[:s.symbolLen] {
638
		if s.norm[i] == notYetAssigned {
639
			var (
640
				end    = tmpTotal + uint64(cnt)*rStep
641
				sStart = uint32(tmpTotal >> vStepLog)
642
				sEnd   = uint32(end >> vStepLog)
643
				weight = sEnd - sStart
644
			)
645
			if weight < 1 {
646
				return errors.New("weight < 1")
647
			}
648
			s.norm[i] = int16(weight)
649
			tmpTotal = end
650
		}
651
	}
652
	return nil
653
}
654

655
// validateNorm validates the normalized histogram table.
656
func (s *Scratch) validateNorm() (err error) {
657
	var total int
658
	for _, v := range s.norm[:s.symbolLen] {
659
		if v >= 0 {
660
			total += int(v)
661
		} else {
662
			total -= int(v)
663
		}
664
	}
665
	defer func() {
666
		if err == nil {
667
			return
668
		}
669
		fmt.Printf("selected TableLog: %d, Symbol length: %d\n", s.actualTableLog, s.symbolLen)
670
		for i, v := range s.norm[:s.symbolLen] {
671
			fmt.Printf("%3d: %5d -> %4d \n", i, s.count[i], v)
672
		}
673
	}()
674
	if total != (1 << s.actualTableLog) {
675
		return fmt.Errorf("warning: Total == %d != %d", total, 1<<s.actualTableLog)
676
	}
677
	for i, v := range s.count[s.symbolLen:] {
678
		if v != 0 {
679
			return fmt.Errorf("warning: Found symbol out of range, %d after cut", i)
680
		}
681
	}
682
	return nil
683
}
684