code wiki / _hdl_build / nx_h264_emit.nx
`nx_h264_emit.nx` source

↩ module page · 249 lines · 9940 B
1// nx_h264_emit.nx -- emit a VALID H.264 Annex-B bitstream: SPS + PPS + IDR slice of I_PCM macroblocks.
2//
3// argv: <outfile> [mbw] [mbh]
4//
5// WHY I_PCM: it carries raw samples, so a decoder that accepts this stream has exercised our NAL
6// framing, emulation-prevention, exp-golomb writer, SPS, PPS and slice header WITHOUT us needing a
7// transform, prediction or entropy coder. It isolates BITSTREAM CONFORMANCE from compression, which
8// is exactly the axis `video|codec-conformance` declares. The pixels come back bit-exact or the
9// stream was wrong -- there is no "close enough" to hide behind.
10//
11// The whole point is that the judge is EXTERNAL: ffmpeg is the reference decoder and it is installed
12// on this host. Our own gates cannot grade this; a decoder we did not write either accepts the
13// stream or does not.
14// license_tier: ORIGINAL   No hw writes (Rule 26).
15import "nx_seg_store.nx"
16
17// Local print/parse helpers: _p/_fn live in the ecomat lib, which this organ has no reason to pull
18// in, and nx_seg_store exports no atoi. Defining them here keeps the emitter dependency-light --
19// it should need a bit writer and a file descriptor, nothing else.
20func _p(s: *u8) -> i64 {
  var n: i64 = 0
  while s[n] != (0 as u8) { n = n + 1 }
  sys_write(1, s, n)
  return 0
25}
26
27func _fn(fd: i64, v: i64) -> i64 {
  let t: *u8 = sys_mmap(32)
  var m: i64 = v
  var neg: i64 = 0
  if m < 0 { neg = 1; m = 0 - m }
  var k: i64 = 0
  if m == 0 { t[0] = 48 as u8; k = 1 }
  while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
  let o: *u8 = sys_mmap(40)
  var p: i64 = 0
  if neg == 1 { o[p] = 45 as u8; p = p + 1 }
  var i: i64 = 0
  while i < k { o[p] = t[k - 1 - i]; p = p + 1; i = i + 1 }
  sys_write(fd, o, p)
  return 0
42}
43
44func h_atoi(s: *u8) -> i64 {
  var v: i64 = 0
  var i: i64 = 0
  while s[i] != (0 as u8) {
      let c: i64 = s[i] as i64
      if c >= 48 { if c <= 57 { v = (v * 10) + (c - 48) } }
      i = i + 1
  }
  return v
53}
54
55const H_MODE: i64 = 0x1a4
56const H_MAXBUF: i64 = 4194304
57
58// ---- bit writer: MSB-first into a byte buffer, pos measured in BITS ----
59func bw_put(buf: *u8, posbox: *i64, val: i64, nbits: i64) -> i64 {
  var i: i64 = nbits - 1
  while i >= 0 {
      let bit: i64 = (val >> i) & 1
      let p: i64 = posbox[0]
      let byi: i64 = p >> 3
      let bii: i64 = 7 - (p & 7)
      if bit == 1 { buf[byi] = ((buf[byi] as i64) | (1 << bii)) as u8 } else { buf[byi] = ((buf[byi] as i64) & (255 - (1 << bii))) as u8 }
      posbox[0] = p + 1
      i = i - 1
  }
  return 0
71}
72
73func bw_bitlen(v: i64) -> i64 {
  var n: i64 = 0
  var x: i64 = v
  while x > 0 { n = n + 1; x = x >> 1 }
  return n
78}
79
80// unsigned exp-golomb: codeNum v -> (n-1) zeros, then (v+1) in n bits, n = bitlen(v+1)
81func bw_ue(buf: *u8, posbox: *i64, v: i64) -> i64 {
  let c: i64 = v + 1
  let n: i64 = bw_bitlen(c)
  bw_put(buf, posbox, 0, n - 1)
  bw_put(buf, posbox, c, n)
  return 0
87}
88
89// signed exp-golomb: k>0 -> 2k-1, k<=0 -> -2k
90func bw_se(buf: *u8, posbox: *i64, v: i64) -> i64 {
  if v > 0 { return bw_ue(buf, posbox, (2 * v) - 1) }
  return bw_ue(buf, posbox, 0 - (2 * v))
93}
94
95// rbsp_trailing_bits: a single 1 then zeros to the byte boundary
96func bw_trailing(buf: *u8, posbox: *i64) -> i64 {
  bw_put(buf, posbox, 1, 1)
  while (posbox[0] & 7) != 0 { bw_put(buf, posbox, 0, 1) }
  return 0
100}
101
102// pcm_alignment_zero_bits: zeros to the byte boundary (no leading 1)
103func bw_align0(buf: *u8, posbox: *i64) -> i64 {
  while (posbox[0] & 7) != 0 { bw_put(buf, posbox, 0, 1) }
  return 0
106}
107
108// ---- NAL assembly: start code + header byte + RBSP with emulation prevention ----
109// A raw 00 00 00/01/02/03 inside the payload would be mistaken for a start code, so the spec
110// requires an escape byte 0x03 after any two consecutive zero bytes. Omitting this is the classic
111// way a stream decodes fine on short test data and corrupts on real data.
112func nal_emit(out: *u8, obox: *i64, nal_ref_idc: i64, nal_type: i64, rbsp: *u8, rlen: i64) -> i64 {
  var o: i64 = obox[0]
  out[o] = 0 as u8; o = o + 1
  out[o] = 0 as u8; o = o + 1
  out[o] = 0 as u8; o = o + 1
  out[o] = 1 as u8; o = o + 1
  out[o] = (((nal_ref_idc & 3) << 5) | (nal_type & 31)) as u8
  o = o + 1
  var zeros: i64 = 0
  var i: i64 = 0
  while i < rlen {
      let b: i64 = rbsp[i] as i64
      if zeros >= 2 { if b <= 3 {
          out[o] = 3 as u8
          o = o + 1
          zeros = 0
      } }
      out[o] = b as u8
      o = o + 1
      if b == 0 { zeros = zeros + 1 } else { zeros = 0 }
      i = i + 1
  }
  obox[0] = o
  return 0
136}
137
138func main(argc: i64, argv: *i64) -> i64 {
  if argc < 2 {
      _p("usage: nx_h264_emit <outfile> [mb_width] [mb_height]\n" as *u8)
      _p("  emits SPS+PPS+IDR(I_PCM) Annex-B; verify with the ffmpeg reference decoder.\n" as *u8)
      sys_exit(1)
      return 1
  }
  let outpath: *u8 = argv[1] as *u8
  var mbw: i64 = 1
  var mbh: i64 = 1
  if argc >= 3 { mbw = h_atoi(argv[2] as *u8) }
  if argc >= 4 { mbh = h_atoi(argv[3] as *u8) }
  if mbw < 1 { mbw = 1 }
  if mbh < 1 { mbh = 1 }
152
  let rbsp: *u8 = sys_mmap(H_MAXBUF)
  let out: *u8 = sys_mmap(H_MAXBUF)
  let pos: *i64 = sys_mmap(16) as *i64
  let obox: *i64 = sys_mmap(16) as *i64
  obox[0] = 0
158
  // ---------------- SPS (nal_type 7) ----------------
  pos[0] = 0
  bw_put(rbsp, pos, 66, 8)          // profile_idc = 66 (Baseline)
  bw_put(rbsp, pos, 0, 8)           // constraint_set flags + reserved_zero_2bits
  bw_put(rbsp, pos, 10, 8)          // level_idc = 1.0
  bw_ue(rbsp, pos, 0)               // seq_parameter_set_id
  bw_ue(rbsp, pos, 0)               // log2_max_frame_num_minus4 -> 4
  bw_ue(rbsp, pos, 2)               // pic_order_cnt_type = 2 (decode order == display order)
  bw_ue(rbsp, pos, 0)               // max_num_ref_frames
  bw_put(rbsp, pos, 0, 1)           // gaps_in_frame_num_value_allowed_flag
  bw_ue(rbsp, pos, mbw - 1)         // pic_width_in_mbs_minus1
  bw_ue(rbsp, pos, mbh - 1)         // pic_height_in_map_units_minus1
  bw_put(rbsp, pos, 1, 1)           // frame_mbs_only_flag
  bw_put(rbsp, pos, 0, 1)           // direct_8x8_inference_flag
  bw_put(rbsp, pos, 0, 1)           // frame_cropping_flag
  bw_put(rbsp, pos, 0, 1)           // vui_parameters_present_flag
  bw_trailing(rbsp, pos)
  nal_emit(out, obox, 3, 7, rbsp, pos[0] >> 3)
177
  // ---------------- PPS (nal_type 8) ----------------
  pos[0] = 0
  bw_ue(rbsp, pos, 0)               // pic_parameter_set_id
  bw_ue(rbsp, pos, 0)               // seq_parameter_set_id
  bw_put(rbsp, pos, 0, 1)           // entropy_coding_mode_flag = 0 (CAVLC)
  bw_put(rbsp, pos, 0, 1)           // bottom_field_pic_order_in_frame_present_flag
  bw_ue(rbsp, pos, 0)               // num_slice_groups_minus1
  bw_ue(rbsp, pos, 0)               // num_ref_idx_l0_default_active_minus1
  bw_ue(rbsp, pos, 0)               // num_ref_idx_l1_default_active_minus1
  bw_put(rbsp, pos, 0, 1)           // weighted_pred_flag
  bw_put(rbsp, pos, 0, 2)           // weighted_bipred_idc
  bw_se(rbsp, pos, 0)               // pic_init_qp_minus26
  bw_se(rbsp, pos, 0)               // pic_init_qs_minus26
  bw_se(rbsp, pos, 0)               // chroma_qp_index_offset
  bw_put(rbsp, pos, 0, 1)           // deblocking_filter_control_present_flag
  bw_put(rbsp, pos, 0, 1)           // constrained_intra_pred_flag
  bw_put(rbsp, pos, 0, 1)           // redundant_pic_cnt_present_flag
  bw_trailing(rbsp, pos)
  nal_emit(out, obox, 3, 8, rbsp, pos[0] >> 3)
197
  // ---------------- IDR slice (nal_type 5) ----------------
  pos[0] = 0
  bw_ue(rbsp, pos, 0)               // first_mb_in_slice
  bw_ue(rbsp, pos, 7)               // slice_type 7 = I (all slices in pic are I)
  bw_ue(rbsp, pos, 0)               // pic_parameter_set_id
  bw_put(rbsp, pos, 0, 4)           // frame_num, u(log2_max_frame_num = 4)
  bw_ue(rbsp, pos, 0)               // idr_pic_id (IdrPicFlag)
  bw_put(rbsp, pos, 0, 1)           // dec_ref_pic_marking: no_output_of_prior_pics_flag
  bw_put(rbsp, pos, 0, 1)           // dec_ref_pic_marking: long_term_reference_flag
  bw_se(rbsp, pos, 0)               // slice_qp_delta
208
  // slice_data: every macroblock is I_PCM (mb_type 25 in an I slice)
  let nmb: i64 = mbw * mbh
  var m: i64 = 0
  while m < nmb {
      bw_ue(rbsp, pos, 25)          // mb_type = I_PCM
      bw_align0(rbsp, pos)          // pcm_alignment_zero_bits
      let mbx: i64 = m % mbw
      let mby: i64 = m / mbw
      // 256 luma samples: a deterministic ramp so a wrong stride shows up as a visible shear
      var y: i64 = 0
      while y < 16 {
          var x: i64 = 0
          while x < 16 {
              bw_put(rbsp, pos, ((y * 16) + x) & 255, 8)
              x = x + 1
          }
          y = y + 1
      }
      // 64 Cb then 64 Cr, varied per macroblock so MB ordering is checkable too
      var c: i64 = 0
      while c < 64 { bw_put(rbsp, pos, (64 + mbx) & 255, 8); c = c + 1 }
      c = 0
      while c < 64 { bw_put(rbsp, pos, (192 - mby) & 255, 8); c = c + 1 }
      m = m + 1
  }
  bw_trailing(rbsp, pos)
  nal_emit(out, obox, 3, 5, rbsp, pos[0] >> 3)
236
  let fd: i64 = sys_openat_wr(outpath, H_MODE)
  if fd < 0 { _p("NX-H264-EMIT RED reason=cannot-open-output\n" as *u8); sys_exit(2); return 2 }
  sys_write(fd, out, obox[0])
  sys_close(fd)
241
  _p("NX-H264-EMIT ok " as *u8); _p(outpath)
  _p("  mbs=" as *u8); _fn(1, nmb)
  _p("  res=" as *u8); _fn(1, mbw * 16); _p("x" as *u8); _fn(1, mbh * 16)
  _p("  bytes=" as *u8); _fn(1, obox[0]); _p("\n" as *u8)
  _p("  verify EXTERNALLY: ffmpeg -i " as *u8); _p(outpath); _p(" -f rawvideo -pix_fmt yuv420p out.yuv\n" as *u8)
  sys_exit(0)
  return 0
249}
nx_h264_emit.nx source

`nx_h264_emit.nx` source
