Pragya Arya | 324337e | 2020-02-20 14:35:08 +0530 | [diff] [blame] | 1 | // Copyright 2019+ Klaus Post. All rights reserved. |
| 2 | // License information can be found in the LICENSE file. |
| 3 | // Based on work by Yann Collet, released under BSD License. |
| 4 | |
| 5 | package zstd |
| 6 | |
| 7 | import ( |
| 8 | "errors" |
| 9 | "fmt" |
| 10 | "math" |
| 11 | "math/bits" |
| 12 | |
| 13 | "github.com/klauspost/compress/huff0" |
| 14 | ) |
| 15 | |
| 16 | type blockEnc struct { |
| 17 | size int |
| 18 | literals []byte |
| 19 | sequences []seq |
| 20 | coders seqCoders |
| 21 | litEnc *huff0.Scratch |
| 22 | wr bitWriter |
| 23 | |
| 24 | extraLits int |
| 25 | last bool |
| 26 | |
| 27 | output []byte |
| 28 | recentOffsets [3]uint32 |
| 29 | prevRecentOffsets [3]uint32 |
| 30 | } |
| 31 | |
| 32 | // init should be used once the block has been created. |
| 33 | // If called more than once, the effect is the same as calling reset. |
| 34 | func (b *blockEnc) init() { |
| 35 | if cap(b.literals) < maxCompressedLiteralSize { |
| 36 | b.literals = make([]byte, 0, maxCompressedLiteralSize) |
| 37 | } |
| 38 | const defSeqs = 200 |
| 39 | b.literals = b.literals[:0] |
| 40 | if cap(b.sequences) < defSeqs { |
| 41 | b.sequences = make([]seq, 0, defSeqs) |
| 42 | } |
| 43 | if cap(b.output) < maxCompressedBlockSize { |
| 44 | b.output = make([]byte, 0, maxCompressedBlockSize) |
| 45 | } |
| 46 | if b.coders.mlEnc == nil { |
| 47 | b.coders.mlEnc = &fseEncoder{} |
| 48 | b.coders.mlPrev = &fseEncoder{} |
| 49 | b.coders.ofEnc = &fseEncoder{} |
| 50 | b.coders.ofPrev = &fseEncoder{} |
| 51 | b.coders.llEnc = &fseEncoder{} |
| 52 | b.coders.llPrev = &fseEncoder{} |
| 53 | } |
| 54 | b.litEnc = &huff0.Scratch{WantLogLess: 4} |
| 55 | b.reset(nil) |
| 56 | } |
| 57 | |
| 58 | // initNewEncode can be used to reset offsets and encoders to the initial state. |
| 59 | func (b *blockEnc) initNewEncode() { |
| 60 | b.recentOffsets = [3]uint32{1, 4, 8} |
| 61 | b.litEnc.Reuse = huff0.ReusePolicyNone |
| 62 | b.coders.setPrev(nil, nil, nil) |
| 63 | } |
| 64 | |
| 65 | // reset will reset the block for a new encode, but in the same stream, |
| 66 | // meaning that state will be carried over, but the block content is reset. |
| 67 | // If a previous block is provided, the recent offsets are carried over. |
| 68 | func (b *blockEnc) reset(prev *blockEnc) { |
| 69 | b.extraLits = 0 |
| 70 | b.literals = b.literals[:0] |
| 71 | b.size = 0 |
| 72 | b.sequences = b.sequences[:0] |
| 73 | b.output = b.output[:0] |
| 74 | b.last = false |
| 75 | if prev != nil { |
| 76 | b.recentOffsets = prev.prevRecentOffsets |
| 77 | } |
| 78 | } |
| 79 | |
| 80 | // reset will reset the block for a new encode, but in the same stream, |
| 81 | // meaning that state will be carried over, but the block content is reset. |
| 82 | // If a previous block is provided, the recent offsets are carried over. |
| 83 | func (b *blockEnc) swapEncoders(prev *blockEnc) { |
| 84 | b.coders.swap(&prev.coders) |
| 85 | b.litEnc, prev.litEnc = prev.litEnc, b.litEnc |
| 86 | } |
| 87 | |
| 88 | // blockHeader contains the information for a block header. |
| 89 | type blockHeader uint32 |
| 90 | |
| 91 | // setLast sets the 'last' indicator on a block. |
| 92 | func (h *blockHeader) setLast(b bool) { |
| 93 | if b { |
| 94 | *h = *h | 1 |
| 95 | } else { |
| 96 | const mask = (1 << 24) - 2 |
| 97 | *h = *h & mask |
| 98 | } |
| 99 | } |
| 100 | |
| 101 | // setSize will store the compressed size of a block. |
| 102 | func (h *blockHeader) setSize(v uint32) { |
| 103 | const mask = 7 |
| 104 | *h = (*h)&mask | blockHeader(v<<3) |
| 105 | } |
| 106 | |
| 107 | // setType sets the block type. |
| 108 | func (h *blockHeader) setType(t blockType) { |
| 109 | const mask = 1 | (((1 << 24) - 1) ^ 7) |
| 110 | *h = (*h & mask) | blockHeader(t<<1) |
| 111 | } |
| 112 | |
| 113 | // appendTo will append the block header to a slice. |
| 114 | func (h blockHeader) appendTo(b []byte) []byte { |
| 115 | return append(b, uint8(h), uint8(h>>8), uint8(h>>16)) |
| 116 | } |
| 117 | |
| 118 | // String returns a string representation of the block. |
| 119 | func (h blockHeader) String() string { |
| 120 | return fmt.Sprintf("Type: %d, Size: %d, Last:%t", (h>>1)&3, h>>3, h&1 == 1) |
| 121 | } |
| 122 | |
| 123 | // literalsHeader contains literals header information. |
| 124 | type literalsHeader uint64 |
| 125 | |
| 126 | // setType can be used to set the type of literal block. |
| 127 | func (h *literalsHeader) setType(t literalsBlockType) { |
| 128 | const mask = math.MaxUint64 - 3 |
| 129 | *h = (*h & mask) | literalsHeader(t) |
| 130 | } |
| 131 | |
| 132 | // setSize can be used to set a single size, for uncompressed and RLE content. |
| 133 | func (h *literalsHeader) setSize(regenLen int) { |
| 134 | inBits := bits.Len32(uint32(regenLen)) |
| 135 | // Only retain 2 bits |
| 136 | const mask = 3 |
| 137 | lh := uint64(*h & mask) |
| 138 | switch { |
| 139 | case inBits < 5: |
| 140 | lh |= (uint64(regenLen) << 3) | (1 << 60) |
| 141 | if debug { |
| 142 | got := int(lh>>3) & 0xff |
| 143 | if got != regenLen { |
| 144 | panic(fmt.Sprint("litRegenSize = ", regenLen, "(want) != ", got, "(got)")) |
| 145 | } |
| 146 | } |
| 147 | case inBits < 12: |
| 148 | lh |= (1 << 2) | (uint64(regenLen) << 4) | (2 << 60) |
| 149 | case inBits < 20: |
| 150 | lh |= (3 << 2) | (uint64(regenLen) << 4) | (3 << 60) |
| 151 | default: |
| 152 | panic(fmt.Errorf("internal error: block too big (%d)", regenLen)) |
| 153 | } |
| 154 | *h = literalsHeader(lh) |
| 155 | } |
| 156 | |
| 157 | // setSizes will set the size of a compressed literals section and the input length. |
| 158 | func (h *literalsHeader) setSizes(compLen, inLen int, single bool) { |
| 159 | compBits, inBits := bits.Len32(uint32(compLen)), bits.Len32(uint32(inLen)) |
| 160 | // Only retain 2 bits |
| 161 | const mask = 3 |
| 162 | lh := uint64(*h & mask) |
| 163 | switch { |
| 164 | case compBits <= 10 && inBits <= 10: |
| 165 | if !single { |
| 166 | lh |= 1 << 2 |
| 167 | } |
| 168 | lh |= (uint64(inLen) << 4) | (uint64(compLen) << (10 + 4)) | (3 << 60) |
| 169 | if debug { |
| 170 | const mmask = (1 << 24) - 1 |
| 171 | n := (lh >> 4) & mmask |
| 172 | if int(n&1023) != inLen { |
| 173 | panic(fmt.Sprint("regensize:", int(n&1023), "!=", inLen, inBits)) |
| 174 | } |
| 175 | if int(n>>10) != compLen { |
| 176 | panic(fmt.Sprint("compsize:", int(n>>10), "!=", compLen, compBits)) |
| 177 | } |
| 178 | } |
| 179 | case compBits <= 14 && inBits <= 14: |
| 180 | lh |= (2 << 2) | (uint64(inLen) << 4) | (uint64(compLen) << (14 + 4)) | (4 << 60) |
| 181 | if single { |
| 182 | panic("single stream used with more than 10 bits length.") |
| 183 | } |
| 184 | case compBits <= 18 && inBits <= 18: |
| 185 | lh |= (3 << 2) | (uint64(inLen) << 4) | (uint64(compLen) << (18 + 4)) | (5 << 60) |
| 186 | if single { |
| 187 | panic("single stream used with more than 10 bits length.") |
| 188 | } |
| 189 | default: |
| 190 | panic("internal error: block too big") |
| 191 | } |
| 192 | *h = literalsHeader(lh) |
| 193 | } |
| 194 | |
| 195 | // appendTo will append the literals header to a byte slice. |
| 196 | func (h literalsHeader) appendTo(b []byte) []byte { |
| 197 | size := uint8(h >> 60) |
| 198 | switch size { |
| 199 | case 1: |
| 200 | b = append(b, uint8(h)) |
| 201 | case 2: |
| 202 | b = append(b, uint8(h), uint8(h>>8)) |
| 203 | case 3: |
| 204 | b = append(b, uint8(h), uint8(h>>8), uint8(h>>16)) |
| 205 | case 4: |
| 206 | b = append(b, uint8(h), uint8(h>>8), uint8(h>>16), uint8(h>>24)) |
| 207 | case 5: |
| 208 | b = append(b, uint8(h), uint8(h>>8), uint8(h>>16), uint8(h>>24), uint8(h>>32)) |
| 209 | default: |
| 210 | panic(fmt.Errorf("internal error: literalsHeader has invalid size (%d)", size)) |
| 211 | } |
| 212 | return b |
| 213 | } |
| 214 | |
| 215 | // size returns the output size with currently set values. |
| 216 | func (h literalsHeader) size() int { |
| 217 | return int(h >> 60) |
| 218 | } |
| 219 | |
| 220 | func (h literalsHeader) String() string { |
| 221 | return fmt.Sprintf("Type: %d, SizeFormat: %d, Size: 0x%d, Bytes:%d", literalsBlockType(h&3), (h>>2)&3, h&((1<<60)-1)>>4, h>>60) |
| 222 | } |
| 223 | |
| 224 | // pushOffsets will push the recent offsets to the backup store. |
| 225 | func (b *blockEnc) pushOffsets() { |
| 226 | b.prevRecentOffsets = b.recentOffsets |
| 227 | } |
| 228 | |
| 229 | // pushOffsets will push the recent offsets to the backup store. |
| 230 | func (b *blockEnc) popOffsets() { |
| 231 | b.recentOffsets = b.prevRecentOffsets |
| 232 | } |
| 233 | |
| 234 | // matchOffset will adjust recent offsets and return the adjusted one, |
| 235 | // if it matches a previous offset. |
| 236 | func (b *blockEnc) matchOffset(offset, lits uint32) uint32 { |
| 237 | // Check if offset is one of the recent offsets. |
| 238 | // Adjusts the output offset accordingly. |
| 239 | // Gives a tiny bit of compression, typically around 1%. |
| 240 | if true { |
| 241 | if lits > 0 { |
| 242 | switch offset { |
| 243 | case b.recentOffsets[0]: |
| 244 | offset = 1 |
| 245 | case b.recentOffsets[1]: |
| 246 | b.recentOffsets[1] = b.recentOffsets[0] |
| 247 | b.recentOffsets[0] = offset |
| 248 | offset = 2 |
| 249 | case b.recentOffsets[2]: |
| 250 | b.recentOffsets[2] = b.recentOffsets[1] |
| 251 | b.recentOffsets[1] = b.recentOffsets[0] |
| 252 | b.recentOffsets[0] = offset |
| 253 | offset = 3 |
| 254 | default: |
| 255 | b.recentOffsets[2] = b.recentOffsets[1] |
| 256 | b.recentOffsets[1] = b.recentOffsets[0] |
| 257 | b.recentOffsets[0] = offset |
| 258 | offset += 3 |
| 259 | } |
| 260 | } else { |
| 261 | switch offset { |
| 262 | case b.recentOffsets[1]: |
| 263 | b.recentOffsets[1] = b.recentOffsets[0] |
| 264 | b.recentOffsets[0] = offset |
| 265 | offset = 1 |
| 266 | case b.recentOffsets[2]: |
| 267 | b.recentOffsets[2] = b.recentOffsets[1] |
| 268 | b.recentOffsets[1] = b.recentOffsets[0] |
| 269 | b.recentOffsets[0] = offset |
| 270 | offset = 2 |
| 271 | case b.recentOffsets[0] - 1: |
| 272 | b.recentOffsets[2] = b.recentOffsets[1] |
| 273 | b.recentOffsets[1] = b.recentOffsets[0] |
| 274 | b.recentOffsets[0] = offset |
| 275 | offset = 3 |
| 276 | default: |
| 277 | b.recentOffsets[2] = b.recentOffsets[1] |
| 278 | b.recentOffsets[1] = b.recentOffsets[0] |
| 279 | b.recentOffsets[0] = offset |
| 280 | offset += 3 |
| 281 | } |
| 282 | } |
| 283 | } else { |
| 284 | offset += 3 |
| 285 | } |
| 286 | return offset |
| 287 | } |
| 288 | |
| 289 | // encodeRaw can be used to set the output to a raw representation of supplied bytes. |
| 290 | func (b *blockEnc) encodeRaw(a []byte) { |
| 291 | var bh blockHeader |
| 292 | bh.setLast(b.last) |
| 293 | bh.setSize(uint32(len(a))) |
| 294 | bh.setType(blockTypeRaw) |
| 295 | b.output = bh.appendTo(b.output[:0]) |
| 296 | b.output = append(b.output, a...) |
| 297 | if debug { |
| 298 | println("Adding RAW block, length", len(a)) |
| 299 | } |
| 300 | } |
| 301 | |
| 302 | // encodeRaw can be used to set the output to a raw representation of supplied bytes. |
| 303 | func (b *blockEnc) encodeRawTo(dst, src []byte) []byte { |
| 304 | var bh blockHeader |
| 305 | bh.setLast(b.last) |
| 306 | bh.setSize(uint32(len(src))) |
| 307 | bh.setType(blockTypeRaw) |
| 308 | dst = bh.appendTo(dst) |
| 309 | dst = append(dst, src...) |
| 310 | if debug { |
| 311 | println("Adding RAW block, length", len(src)) |
| 312 | } |
| 313 | return dst |
| 314 | } |
| 315 | |
| 316 | // encodeLits can be used if the block is only litLen. |
| 317 | func (b *blockEnc) encodeLits(raw bool) error { |
| 318 | var bh blockHeader |
| 319 | bh.setLast(b.last) |
| 320 | bh.setSize(uint32(len(b.literals))) |
| 321 | |
| 322 | // Don't compress extremely small blocks |
| 323 | if len(b.literals) < 32 || raw { |
| 324 | if debug { |
| 325 | println("Adding RAW block, length", len(b.literals)) |
| 326 | } |
| 327 | bh.setType(blockTypeRaw) |
| 328 | b.output = bh.appendTo(b.output) |
| 329 | b.output = append(b.output, b.literals...) |
| 330 | return nil |
| 331 | } |
| 332 | |
| 333 | var ( |
| 334 | out []byte |
| 335 | reUsed, single bool |
| 336 | err error |
| 337 | ) |
| 338 | if len(b.literals) >= 1024 { |
| 339 | // Use 4 Streams. |
| 340 | out, reUsed, err = huff0.Compress4X(b.literals, b.litEnc) |
| 341 | } else if len(b.literals) > 32 { |
| 342 | // Use 1 stream |
| 343 | single = true |
| 344 | out, reUsed, err = huff0.Compress1X(b.literals, b.litEnc) |
| 345 | } else { |
| 346 | err = huff0.ErrIncompressible |
| 347 | } |
| 348 | |
| 349 | switch err { |
| 350 | case huff0.ErrIncompressible: |
| 351 | if debug { |
| 352 | println("Adding RAW block, length", len(b.literals)) |
| 353 | } |
| 354 | bh.setType(blockTypeRaw) |
| 355 | b.output = bh.appendTo(b.output) |
| 356 | b.output = append(b.output, b.literals...) |
| 357 | return nil |
| 358 | case huff0.ErrUseRLE: |
| 359 | if debug { |
| 360 | println("Adding RLE block, length", len(b.literals)) |
| 361 | } |
| 362 | bh.setType(blockTypeRLE) |
| 363 | b.output = bh.appendTo(b.output) |
| 364 | b.output = append(b.output, b.literals[0]) |
| 365 | return nil |
| 366 | default: |
| 367 | return err |
| 368 | case nil: |
| 369 | } |
| 370 | // Compressed... |
| 371 | // Now, allow reuse |
| 372 | b.litEnc.Reuse = huff0.ReusePolicyAllow |
| 373 | bh.setType(blockTypeCompressed) |
| 374 | var lh literalsHeader |
| 375 | if reUsed { |
| 376 | if debug { |
| 377 | println("Reused tree, compressed to", len(out)) |
| 378 | } |
| 379 | lh.setType(literalsBlockTreeless) |
| 380 | } else { |
| 381 | if debug { |
| 382 | println("New tree, compressed to", len(out), "tree size:", len(b.litEnc.OutTable)) |
| 383 | } |
| 384 | lh.setType(literalsBlockCompressed) |
| 385 | } |
| 386 | // Set sizes |
| 387 | lh.setSizes(len(out), len(b.literals), single) |
| 388 | bh.setSize(uint32(len(out) + lh.size() + 1)) |
| 389 | |
| 390 | // Write block headers. |
| 391 | b.output = bh.appendTo(b.output) |
| 392 | b.output = lh.appendTo(b.output) |
| 393 | // Add compressed data. |
| 394 | b.output = append(b.output, out...) |
| 395 | // No sequences. |
| 396 | b.output = append(b.output, 0) |
| 397 | return nil |
| 398 | } |
| 399 | |
| 400 | // fuzzFseEncoder can be used to fuzz the FSE encoder. |
| 401 | func fuzzFseEncoder(data []byte) int { |
| 402 | if len(data) > maxSequences || len(data) < 2 { |
| 403 | return 0 |
| 404 | } |
| 405 | enc := fseEncoder{} |
| 406 | hist := enc.Histogram()[:256] |
| 407 | maxSym := uint8(0) |
| 408 | for i, v := range data { |
| 409 | v = v & 63 |
| 410 | data[i] = v |
| 411 | hist[v]++ |
| 412 | if v > maxSym { |
| 413 | maxSym = v |
| 414 | } |
| 415 | } |
| 416 | if maxSym == 0 { |
| 417 | // All 0 |
| 418 | return 0 |
| 419 | } |
| 420 | maxCount := func(a []uint32) int { |
| 421 | var max uint32 |
| 422 | for _, v := range a { |
| 423 | if v > max { |
| 424 | max = v |
| 425 | } |
| 426 | } |
| 427 | return int(max) |
| 428 | } |
| 429 | cnt := maxCount(hist[:maxSym]) |
| 430 | if cnt == len(data) { |
| 431 | // RLE |
| 432 | return 0 |
| 433 | } |
| 434 | enc.HistogramFinished(maxSym, cnt) |
| 435 | err := enc.normalizeCount(len(data)) |
| 436 | if err != nil { |
| 437 | return 0 |
| 438 | } |
| 439 | _, err = enc.writeCount(nil) |
| 440 | if err != nil { |
| 441 | panic(err) |
| 442 | } |
| 443 | return 1 |
| 444 | } |
| 445 | |
| 446 | // encode will encode the block and append the output in b.output. |
| 447 | func (b *blockEnc) encode(raw bool) error { |
| 448 | if len(b.sequences) == 0 { |
| 449 | return b.encodeLits(raw) |
| 450 | } |
| 451 | // We want some difference |
| 452 | if len(b.literals) > (b.size - (b.size >> 5)) { |
| 453 | return errIncompressible |
| 454 | } |
| 455 | |
| 456 | var bh blockHeader |
| 457 | var lh literalsHeader |
| 458 | bh.setLast(b.last) |
| 459 | bh.setType(blockTypeCompressed) |
| 460 | // Store offset of the block header. Needed when we know the size. |
| 461 | bhOffset := len(b.output) |
| 462 | b.output = bh.appendTo(b.output) |
| 463 | |
| 464 | var ( |
| 465 | out []byte |
| 466 | reUsed, single bool |
| 467 | err error |
| 468 | ) |
| 469 | if len(b.literals) >= 1024 && !raw { |
| 470 | // Use 4 Streams. |
| 471 | out, reUsed, err = huff0.Compress4X(b.literals, b.litEnc) |
| 472 | } else if len(b.literals) > 32 && !raw { |
| 473 | // Use 1 stream |
| 474 | single = true |
| 475 | out, reUsed, err = huff0.Compress1X(b.literals, b.litEnc) |
| 476 | } else { |
| 477 | err = huff0.ErrIncompressible |
| 478 | } |
| 479 | |
| 480 | switch err { |
| 481 | case huff0.ErrIncompressible: |
| 482 | lh.setType(literalsBlockRaw) |
| 483 | lh.setSize(len(b.literals)) |
| 484 | b.output = lh.appendTo(b.output) |
| 485 | b.output = append(b.output, b.literals...) |
| 486 | if debug { |
| 487 | println("Adding literals RAW, length", len(b.literals)) |
| 488 | } |
| 489 | case huff0.ErrUseRLE: |
| 490 | lh.setType(literalsBlockRLE) |
| 491 | lh.setSize(len(b.literals)) |
| 492 | b.output = lh.appendTo(b.output) |
| 493 | b.output = append(b.output, b.literals[0]) |
| 494 | if debug { |
| 495 | println("Adding literals RLE") |
| 496 | } |
| 497 | default: |
| 498 | if debug { |
| 499 | println("Adding literals ERROR:", err) |
| 500 | } |
| 501 | return err |
| 502 | case nil: |
| 503 | // Compressed litLen... |
| 504 | if reUsed { |
| 505 | if debug { |
| 506 | println("reused tree") |
| 507 | } |
| 508 | lh.setType(literalsBlockTreeless) |
| 509 | } else { |
| 510 | if debug { |
| 511 | println("new tree, size:", len(b.litEnc.OutTable)) |
| 512 | } |
| 513 | lh.setType(literalsBlockCompressed) |
| 514 | if debug { |
| 515 | _, _, err := huff0.ReadTable(out, nil) |
| 516 | if err != nil { |
| 517 | panic(err) |
| 518 | } |
| 519 | } |
| 520 | } |
| 521 | lh.setSizes(len(out), len(b.literals), single) |
| 522 | if debug { |
| 523 | printf("Compressed %d literals to %d bytes", len(b.literals), len(out)) |
| 524 | println("Adding literal header:", lh) |
| 525 | } |
| 526 | b.output = lh.appendTo(b.output) |
| 527 | b.output = append(b.output, out...) |
| 528 | b.litEnc.Reuse = huff0.ReusePolicyAllow |
| 529 | if debug { |
| 530 | println("Adding literals compressed") |
| 531 | } |
| 532 | } |
| 533 | // Sequence compression |
| 534 | |
| 535 | // Write the number of sequences |
| 536 | switch { |
| 537 | case len(b.sequences) < 128: |
| 538 | b.output = append(b.output, uint8(len(b.sequences))) |
| 539 | case len(b.sequences) < 0x7f00: // TODO: this could be wrong |
| 540 | n := len(b.sequences) |
| 541 | b.output = append(b.output, 128+uint8(n>>8), uint8(n)) |
| 542 | default: |
| 543 | n := len(b.sequences) - 0x7f00 |
| 544 | b.output = append(b.output, 255, uint8(n), uint8(n>>8)) |
| 545 | } |
| 546 | if debug { |
| 547 | println("Encoding", len(b.sequences), "sequences") |
| 548 | } |
| 549 | b.genCodes() |
| 550 | llEnc := b.coders.llEnc |
| 551 | ofEnc := b.coders.ofEnc |
| 552 | mlEnc := b.coders.mlEnc |
| 553 | err = llEnc.normalizeCount(len(b.sequences)) |
| 554 | if err != nil { |
| 555 | return err |
| 556 | } |
| 557 | err = ofEnc.normalizeCount(len(b.sequences)) |
| 558 | if err != nil { |
| 559 | return err |
| 560 | } |
| 561 | err = mlEnc.normalizeCount(len(b.sequences)) |
| 562 | if err != nil { |
| 563 | return err |
| 564 | } |
| 565 | |
| 566 | // Choose the best compression mode for each type. |
| 567 | // Will evaluate the new vs predefined and previous. |
| 568 | chooseComp := func(cur, prev, preDef *fseEncoder) (*fseEncoder, seqCompMode) { |
| 569 | // See if predefined/previous is better |
| 570 | hist := cur.count[:cur.symbolLen] |
| 571 | nSize := cur.approxSize(hist) + cur.maxHeaderSize() |
| 572 | predefSize := preDef.approxSize(hist) |
| 573 | prevSize := prev.approxSize(hist) |
| 574 | |
| 575 | // Add a small penalty for new encoders. |
| 576 | // Don't bother with extremely small (<2 byte gains). |
| 577 | nSize = nSize + (nSize+2*8*16)>>4 |
| 578 | switch { |
| 579 | case predefSize <= prevSize && predefSize <= nSize || forcePreDef: |
| 580 | if debug { |
| 581 | println("Using predefined", predefSize>>3, "<=", nSize>>3) |
| 582 | } |
| 583 | return preDef, compModePredefined |
| 584 | case prevSize <= nSize: |
| 585 | if debug { |
| 586 | println("Using previous", prevSize>>3, "<=", nSize>>3) |
| 587 | } |
| 588 | return prev, compModeRepeat |
| 589 | default: |
| 590 | if debug { |
| 591 | println("Using new, predef", predefSize>>3, ". previous:", prevSize>>3, ">", nSize>>3, "header max:", cur.maxHeaderSize()>>3, "bytes") |
| 592 | println("tl:", cur.actualTableLog, "symbolLen:", cur.symbolLen, "norm:", cur.norm[:cur.symbolLen], "hist", cur.count[:cur.symbolLen]) |
| 593 | } |
| 594 | return cur, compModeFSE |
| 595 | } |
| 596 | } |
| 597 | |
| 598 | // Write compression mode |
| 599 | var mode uint8 |
| 600 | if llEnc.useRLE { |
| 601 | mode |= uint8(compModeRLE) << 6 |
| 602 | llEnc.setRLE(b.sequences[0].llCode) |
| 603 | if debug { |
| 604 | println("llEnc.useRLE") |
| 605 | } |
| 606 | } else { |
| 607 | var m seqCompMode |
| 608 | llEnc, m = chooseComp(llEnc, b.coders.llPrev, &fsePredefEnc[tableLiteralLengths]) |
| 609 | mode |= uint8(m) << 6 |
| 610 | } |
| 611 | if ofEnc.useRLE { |
| 612 | mode |= uint8(compModeRLE) << 4 |
| 613 | ofEnc.setRLE(b.sequences[0].ofCode) |
| 614 | if debug { |
| 615 | println("ofEnc.useRLE") |
| 616 | } |
| 617 | } else { |
| 618 | var m seqCompMode |
| 619 | ofEnc, m = chooseComp(ofEnc, b.coders.ofPrev, &fsePredefEnc[tableOffsets]) |
| 620 | mode |= uint8(m) << 4 |
| 621 | } |
| 622 | |
| 623 | if mlEnc.useRLE { |
| 624 | mode |= uint8(compModeRLE) << 2 |
| 625 | mlEnc.setRLE(b.sequences[0].mlCode) |
| 626 | if debug { |
| 627 | println("mlEnc.useRLE, code: ", b.sequences[0].mlCode, "value", b.sequences[0].matchLen) |
| 628 | } |
| 629 | } else { |
| 630 | var m seqCompMode |
| 631 | mlEnc, m = chooseComp(mlEnc, b.coders.mlPrev, &fsePredefEnc[tableMatchLengths]) |
| 632 | mode |= uint8(m) << 2 |
| 633 | } |
| 634 | b.output = append(b.output, mode) |
| 635 | if debug { |
| 636 | printf("Compression modes: 0b%b", mode) |
| 637 | } |
| 638 | b.output, err = llEnc.writeCount(b.output) |
| 639 | if err != nil { |
| 640 | return err |
| 641 | } |
| 642 | start := len(b.output) |
| 643 | b.output, err = ofEnc.writeCount(b.output) |
| 644 | if err != nil { |
| 645 | return err |
| 646 | } |
| 647 | if false { |
| 648 | println("block:", b.output[start:], "tablelog", ofEnc.actualTableLog, "maxcount:", ofEnc.maxCount) |
| 649 | fmt.Printf("selected TableLog: %d, Symbol length: %d\n", ofEnc.actualTableLog, ofEnc.symbolLen) |
| 650 | for i, v := range ofEnc.norm[:ofEnc.symbolLen] { |
| 651 | fmt.Printf("%3d: %5d -> %4d \n", i, ofEnc.count[i], v) |
| 652 | } |
| 653 | } |
| 654 | b.output, err = mlEnc.writeCount(b.output) |
| 655 | if err != nil { |
| 656 | return err |
| 657 | } |
| 658 | |
| 659 | // Maybe in block? |
| 660 | wr := &b.wr |
| 661 | wr.reset(b.output) |
| 662 | |
| 663 | var ll, of, ml cState |
| 664 | |
| 665 | // Current sequence |
| 666 | seq := len(b.sequences) - 1 |
| 667 | s := b.sequences[seq] |
| 668 | llEnc.setBits(llBitsTable[:]) |
| 669 | mlEnc.setBits(mlBitsTable[:]) |
| 670 | ofEnc.setBits(nil) |
| 671 | |
| 672 | llTT, ofTT, mlTT := llEnc.ct.symbolTT[:256], ofEnc.ct.symbolTT[:256], mlEnc.ct.symbolTT[:256] |
| 673 | |
| 674 | // We have 3 bounds checks here (and in the loop). |
| 675 | // Since we are iterating backwards it is kinda hard to avoid. |
| 676 | llB, ofB, mlB := llTT[s.llCode], ofTT[s.ofCode], mlTT[s.mlCode] |
| 677 | ll.init(wr, &llEnc.ct, llB) |
| 678 | of.init(wr, &ofEnc.ct, ofB) |
| 679 | wr.flush32() |
| 680 | ml.init(wr, &mlEnc.ct, mlB) |
| 681 | |
| 682 | // Each of these lookups also generates a bounds check. |
| 683 | wr.addBits32NC(s.litLen, llB.outBits) |
| 684 | wr.addBits32NC(s.matchLen, mlB.outBits) |
| 685 | wr.flush32() |
| 686 | wr.addBits32NC(s.offset, ofB.outBits) |
| 687 | if debugSequences { |
| 688 | println("Encoded seq", seq, s, "codes:", s.llCode, s.mlCode, s.ofCode, "states:", ll.state, ml.state, of.state, "bits:", llB, mlB, ofB) |
| 689 | } |
| 690 | seq-- |
| 691 | if llEnc.maxBits+mlEnc.maxBits+ofEnc.maxBits <= 32 { |
| 692 | // No need to flush (common) |
| 693 | for seq >= 0 { |
| 694 | s = b.sequences[seq] |
| 695 | wr.flush32() |
| 696 | llB, ofB, mlB := llTT[s.llCode], ofTT[s.ofCode], mlTT[s.mlCode] |
| 697 | // tabelog max is 8 for all. |
| 698 | of.encode(ofB) |
| 699 | ml.encode(mlB) |
| 700 | ll.encode(llB) |
| 701 | wr.flush32() |
| 702 | |
| 703 | // We checked that all can stay within 32 bits |
| 704 | wr.addBits32NC(s.litLen, llB.outBits) |
| 705 | wr.addBits32NC(s.matchLen, mlB.outBits) |
| 706 | wr.addBits32NC(s.offset, ofB.outBits) |
| 707 | |
| 708 | if debugSequences { |
| 709 | println("Encoded seq", seq, s) |
| 710 | } |
| 711 | |
| 712 | seq-- |
| 713 | } |
| 714 | } else { |
| 715 | for seq >= 0 { |
| 716 | s = b.sequences[seq] |
| 717 | wr.flush32() |
| 718 | llB, ofB, mlB := llTT[s.llCode], ofTT[s.ofCode], mlTT[s.mlCode] |
| 719 | // tabelog max is below 8 for each. |
| 720 | of.encode(ofB) |
| 721 | ml.encode(mlB) |
| 722 | ll.encode(llB) |
| 723 | wr.flush32() |
| 724 | |
| 725 | // ml+ll = max 32 bits total |
| 726 | wr.addBits32NC(s.litLen, llB.outBits) |
| 727 | wr.addBits32NC(s.matchLen, mlB.outBits) |
| 728 | wr.flush32() |
| 729 | wr.addBits32NC(s.offset, ofB.outBits) |
| 730 | |
| 731 | if debugSequences { |
| 732 | println("Encoded seq", seq, s) |
| 733 | } |
| 734 | |
| 735 | seq-- |
| 736 | } |
| 737 | } |
| 738 | ml.flush(mlEnc.actualTableLog) |
| 739 | of.flush(ofEnc.actualTableLog) |
| 740 | ll.flush(llEnc.actualTableLog) |
| 741 | err = wr.close() |
| 742 | if err != nil { |
| 743 | return err |
| 744 | } |
| 745 | b.output = wr.out |
| 746 | |
| 747 | if len(b.output)-3-bhOffset >= b.size { |
| 748 | // Maybe even add a bigger margin. |
| 749 | b.litEnc.Reuse = huff0.ReusePolicyNone |
| 750 | return errIncompressible |
| 751 | } |
| 752 | |
| 753 | // Size is output minus block header. |
| 754 | bh.setSize(uint32(len(b.output)-bhOffset) - 3) |
| 755 | if debug { |
| 756 | println("Rewriting block header", bh) |
| 757 | } |
| 758 | _ = bh.appendTo(b.output[bhOffset:bhOffset]) |
| 759 | b.coders.setPrev(llEnc, mlEnc, ofEnc) |
| 760 | return nil |
| 761 | } |
| 762 | |
| 763 | var errIncompressible = errors.New("incompressible") |
| 764 | |
| 765 | func (b *blockEnc) genCodes() { |
| 766 | if len(b.sequences) == 0 { |
| 767 | // nothing to do |
| 768 | return |
| 769 | } |
| 770 | |
| 771 | if len(b.sequences) > math.MaxUint16 { |
| 772 | panic("can only encode up to 64K sequences") |
| 773 | } |
| 774 | // No bounds checks after here: |
| 775 | llH := b.coders.llEnc.Histogram()[:256] |
| 776 | ofH := b.coders.ofEnc.Histogram()[:256] |
| 777 | mlH := b.coders.mlEnc.Histogram()[:256] |
| 778 | for i := range llH { |
| 779 | llH[i] = 0 |
| 780 | } |
| 781 | for i := range ofH { |
| 782 | ofH[i] = 0 |
| 783 | } |
| 784 | for i := range mlH { |
| 785 | mlH[i] = 0 |
| 786 | } |
| 787 | |
| 788 | var llMax, ofMax, mlMax uint8 |
| 789 | for i, seq := range b.sequences { |
| 790 | v := llCode(seq.litLen) |
| 791 | seq.llCode = v |
| 792 | llH[v]++ |
| 793 | if v > llMax { |
| 794 | llMax = v |
| 795 | } |
| 796 | |
| 797 | v = ofCode(seq.offset) |
| 798 | seq.ofCode = v |
| 799 | ofH[v]++ |
| 800 | if v > ofMax { |
| 801 | ofMax = v |
| 802 | } |
| 803 | |
| 804 | v = mlCode(seq.matchLen) |
| 805 | seq.mlCode = v |
| 806 | mlH[v]++ |
| 807 | if v > mlMax { |
| 808 | mlMax = v |
| 809 | if debug && mlMax > maxMatchLengthSymbol { |
| 810 | panic(fmt.Errorf("mlMax > maxMatchLengthSymbol (%d), matchlen: %d", mlMax, seq.matchLen)) |
| 811 | } |
| 812 | } |
| 813 | b.sequences[i] = seq |
| 814 | } |
| 815 | maxCount := func(a []uint32) int { |
| 816 | var max uint32 |
| 817 | for _, v := range a { |
| 818 | if v > max { |
| 819 | max = v |
| 820 | } |
| 821 | } |
| 822 | return int(max) |
| 823 | } |
| 824 | if mlMax > maxMatchLengthSymbol { |
| 825 | panic(fmt.Errorf("mlMax > maxMatchLengthSymbol (%d)", mlMax)) |
| 826 | } |
| 827 | if ofMax > maxOffsetBits { |
| 828 | panic(fmt.Errorf("ofMax > maxOffsetBits (%d)", ofMax)) |
| 829 | } |
| 830 | if llMax > maxLiteralLengthSymbol { |
| 831 | panic(fmt.Errorf("llMax > maxLiteralLengthSymbol (%d)", llMax)) |
| 832 | } |
| 833 | |
| 834 | b.coders.mlEnc.HistogramFinished(mlMax, maxCount(mlH[:mlMax+1])) |
| 835 | b.coders.ofEnc.HistogramFinished(ofMax, maxCount(ofH[:ofMax+1])) |
| 836 | b.coders.llEnc.HistogramFinished(llMax, maxCount(llH[:llMax+1])) |
| 837 | } |