media: vidtv: rewrite the adaption field logic

When compared with a stream generated via ffmpeg, it can be
noticed that the PES doesn't contain any PCR info. That could
cause problems with userspace decoding. So, rewrite the
logic that fills the adaptation info, in order to allow it
to add PCR frames without breaking frame alignment.

Signed-off-by: Mauro Carvalho Chehab <mchehab+huawei@kernel.org>

3 files changed, 126 insertions, 71 deletions

diff --git a/drivers/media/test-drivers/vidtv/vidtv_mux.c b/drivers/media/test-drivers/vidtv/vidtv_mux.c
index 6a3123c51cc4..082740ae9d44 100644
--- a/drivers/media/test-drivers/vidtv/vidtv_mux.c
+++ b/drivers/media/test-drivers/vidtv/vidtv_mux.c
@@ -254,6 +254,7 @@ static u32 vidtv_mux_packetize_access_units(struct vidtv_mux *m,
 		args.access_unit_len = au->nbytes;
 		args.dest_offset     = m->mux_buf_offset;
 		args.pts             = au->pts;
+		args.pcr	     = m->timing.clk;
 
 		m->mux_buf_offset += vidtv_pes_write_into(args);
 
diff --git a/drivers/media/test-drivers/vidtv/vidtv_pes.c b/drivers/media/test-drivers/vidtv/vidtv_pes.c
index de0ce5529d06..cbb3fceb00ad 100644
--- a/drivers/media/test-drivers/vidtv/vidtv_pes.c
+++ b/drivers/media/test-drivers/vidtv/vidtv_pes.c
@@ -43,11 +43,14 @@ static u32 vidtv_pes_op_get_len(bool send_pts, bool send_dts)
 	return len;
 }
 
+#define SIZE_PCR (6 + sizeof(struct vidtv_mpeg_ts_adaption))
+
 static u32 vidtv_pes_h_get_len(bool send_pts, bool send_dts)
 {
-	/* PES header length notwithstanding stuffing bytes */
 	u32 len = 0;
 
+	/* PES header length notwithstanding stuffing bytes */
+
 	len += sizeof(struct vidtv_mpeg_pes);
 	len += vidtv_pes_op_get_len(send_pts, send_dts);
 
@@ -176,39 +179,58 @@ static u32 vidtv_pes_write_h(struct pes_header_write_args args)
 	return nbytes;
 }
 
+static u32 vidtv_pes_write_pcr_bits(u8 *to, u32 to_offset, u64 pcr)
+{
+	/* Exact same from ffmpeg. PCR is a counter driven by a 27Mhz clock */
+	u64 div;
+	u64 rem;
+	u8 *buf = to + to_offset;
+	u64 pcr_low;
+	u64 pcr_high;
+
+	div = div64_u64_rem(pcr, 300, &rem);
+
+	pcr_low = rem; /* pcr_low = pcr % 300 */
+	pcr_high = div; /* pcr_high = pcr / 300 */
+
+	*buf++ = pcr_high >> 25;
+	*buf++ = pcr_high >> 17;
+	*buf++ = pcr_high >>  9;
+	*buf++ = pcr_high >>  1;
+	*buf++ = pcr_high <<  7 | pcr_low >> 8 | 0x7e;
+	*buf++ = pcr_low;
+
+	return 6;
+}
+
 static u32 vidtv_pes_write_stuffing(struct pes_ts_header_write_args *args,
-				    u32 dest_offset)
+				    u32 dest_offset, bool need_pcr,
+				    u64 *last_pcr)
 {
-	u32 nbytes = 0;
 	struct vidtv_mpeg_ts_adaption ts_adap = {};
-	u32 stuff_nbytes;
+	int stuff_nbytes;
+	u32 nbytes = 0;
 
 	if (!args->n_stuffing_bytes)
-		goto out;
+		return 0;
 
-	if (args->n_stuffing_bytes > PES_TS_HEADER_MAX_STUFFING_BYTES) {
-		pr_warn_ratelimited("More than %d stuffing bytes for a PES packet!\n",
-				    PES_TS_HEADER_MAX_STUFFING_BYTES);
+	ts_adap.random_access = 1;
 
-		args->n_stuffing_bytes = PES_TS_HEADER_MAX_STUFFING_BYTES;
+	/* length _immediately_ following 'adaptation_field_length' */
+	if (need_pcr) {
+		ts_adap.PCR = 1;
+		ts_adap.length = SIZE_PCR;
+	} else {
+		ts_adap.length = sizeof(ts_adap);
 	}
+	stuff_nbytes = args->n_stuffing_bytes - ts_adap.length;
 
-	/* the AF will only be its 'length' field with a value of zero */
-	if (args->n_stuffing_bytes == 1) {
-		nbytes += vidtv_memset(args->dest_buf,
-				       dest_offset + nbytes,
-				       args->dest_buf_sz,
-				       0,
-				       args->n_stuffing_bytes);
-		goto out;
-	}
+	ts_adap.length -= sizeof(ts_adap.length);
 
-	stuff_nbytes = args->n_stuffing_bytes - sizeof(ts_adap);
+	if (unlikely(stuff_nbytes < 0))
+		stuff_nbytes = 0;
 
-	/* length _immediately_ following 'adaptation_field_length' */
-	ts_adap.length = sizeof(ts_adap) -
-			 sizeof(ts_adap.length) +
-			 stuff_nbytes;
+	ts_adap.length += stuff_nbytes;
 
 	/* write the adap after the TS header */
 	nbytes += vidtv_memcpy(args->dest_buf,
@@ -217,23 +239,37 @@ static u32 vidtv_pes_write_stuffing(struct pes_ts_header_write_args *args,
 			       &ts_adap,
 			       sizeof(ts_adap));
 
-	/* write the stuffing bytes */
-	nbytes += vidtv_memset(args->dest_buf,
-			       dest_offset + nbytes,
-			       args->dest_buf_sz,
-			       TS_FILL_BYTE,
-			       stuff_nbytes);
+	/* write the optional PCR */
+	if (need_pcr) {
+		nbytes += vidtv_pes_write_pcr_bits(args->dest_buf,
+						   dest_offset + nbytes,
+						   args->pcr);
+
+		*last_pcr = args->pcr;
+	}
 
-out:
+	/* write the stuffing bytes, if are there anything left */
+	if (stuff_nbytes)
+		nbytes += vidtv_memset(args->dest_buf,
+				       dest_offset + nbytes,
+				       args->dest_buf_sz,
+				       TS_FILL_BYTE,
+				       stuff_nbytes);
+
+	/*
+	 * The n_stuffing_bytes contain a pre-calculated value of
+	 * the amount of data that this function would read, made from
+	 * vidtv_pes_h_get_len(). If something went wrong, print a warning
+	 */
 	if (nbytes != args->n_stuffing_bytes)
 		pr_warn_ratelimited("write size was %d, expected %d\n",
-				    nbytes,
-				    args->n_stuffing_bytes);
+				    nbytes, args->n_stuffing_bytes);
 
 	return nbytes;
 }
 
-static u32 vidtv_pes_write_ts_h(struct pes_ts_header_write_args args)
+static u32 vidtv_pes_write_ts_h(struct pes_ts_header_write_args args,
+				bool need_pcr, u64 *last_pcr)
 {
 	/* number of bytes written by this function */
 	u32 nbytes = 0;
@@ -258,43 +294,35 @@ static u32 vidtv_pes_write_ts_h(struct pes_ts_header_write_args args)
 			       sizeof(ts_header));
 
 	/* write stuffing, if any */
-	nbytes += vidtv_pes_write_stuffing(&args, args.dest_offset + nbytes);
+	nbytes += vidtv_pes_write_stuffing(&args, args.dest_offset + nbytes,
+					   need_pcr, last_pcr);
 
 	return nbytes;
 }
 
 u32 vidtv_pes_write_into(struct pes_write_args args)
 {
-	u32 nbytes_past_boundary = (args.dest_offset % TS_PACKET_LEN);
-	bool aligned = (nbytes_past_boundary == 0);
-
+	u32 unaligned_bytes = (args.dest_offset % TS_PACKET_LEN);
 	struct pes_ts_header_write_args ts_header_args = {};
-	struct pes_header_write_args pes_header_args   = {};
-
-	/* number of bytes written by this function */
-	u32 nbytes        = 0;
+	struct pes_header_write_args pes_header_args = {};
 	u32 remaining_len = args.access_unit_len;
-
 	bool wrote_pes_header = false;
+	u64 last_pcr = args.pcr;
+	bool need_pcr = true;
+	u32 available_space;
+	u32 payload_size;
+	u32 stuff_bytes;
+	u32 nbytes = 0;
 
-	/* whether we need to stuff the TS packet to align the buffer */
-	bool should_insert_stuffing_bytes = false;
-
-	u32 available_space    = 0;
-	u32 payload_write_len  = 0;
-	u32 num_stuffing_bytes = 0;
-
-	if (!aligned) {
-		pr_warn_ratelimited("Cannot start a PES packet in a misaligned buffer\n");
+	if (unaligned_bytes) {
+		pr_warn_ratelimited("buffer is misaligned, while starting PES\n");
 
 		/* forcibly align and hope for the best */
 		nbytes += vidtv_memset(args.dest_buf,
 				       args.dest_offset + nbytes,
 				       args.dest_buf_sz,
 				       TS_FILL_BYTE,
-				       TS_PACKET_LEN - nbytes_past_boundary);
-
-		aligned = true;
+				       TS_PACKET_LEN - unaligned_bytes);
 	}
 
 	if (args.send_dts && !args.send_pts) {
@@ -310,44 +338,68 @@ u32 vidtv_pes_write_into(struct pes_write_args args)
 	}
 
 	while (remaining_len) {
+		available_space = TS_PAYLOAD_LEN;
 		/*
 		 * The amount of space initially available in the TS packet.
 		 * if this is the beginning of the PES packet, take into account
 		 * the space needed for the TS header _and_ for the PES header
 		 */
-		available_space = (!wrote_pes_header) ?
-				  TS_PAYLOAD_LEN -
-				  vidtv_pes_h_get_len(args.send_pts, args.send_dts) :
-				  TS_PAYLOAD_LEN;
+		if (!wrote_pes_header)
+			available_space -= vidtv_pes_h_get_len(args.send_pts,
+							       args.send_dts);
 
-		/* if the encoder has inserted stuffing bytes in the PES
+		/*
+		 * if the encoder has inserted stuffing bytes in the PES
 		 * header, account for them.
 		 */
 		available_space -= args.n_pes_h_s_bytes;
 
-		/* whether stuffing bytes are needed to align the buffer */
-		should_insert_stuffing_bytes = remaining_len < available_space;
+		/* Take the extra adaptation into account if need to send PCR */
+		if (need_pcr) {
+			available_space -= SIZE_PCR;
+			stuff_bytes = SIZE_PCR;
+		} else {
+			stuff_bytes = 0;
+		}
 
 		/*
 		 * how much of the _actual_ payload should be written in this
 		 * packet.
 		 */
-		payload_write_len = (should_insert_stuffing_bytes) ?
-				    remaining_len :
-				    available_space;
-
-		num_stuffing_bytes = available_space - payload_write_len;
+		if (remaining_len >= available_space) {
+			payload_size = available_space;
+		} else {
+			/* Last frame should ensure 188-bytes PS alignment */
+			payload_size = remaining_len;
+			stuff_bytes += available_space - payload_size;
+
+			/*
+			 * Ensure that the stuff bytes will be within the
+			 * allowed range, decrementing the number of payload
+			 * bytes to write if needed.
+			 */
+			if (stuff_bytes > PES_TS_HEADER_MAX_STUFFING_BYTES) {
+				u32 tmp = stuff_bytes - PES_TS_HEADER_MAX_STUFFING_BYTES;
+
+				stuff_bytes = PES_TS_HEADER_MAX_STUFFING_BYTES;
+				payload_size -= tmp;
+			}
+		}
 
 		/* write ts header */
 		ts_header_args.dest_buf           = args.dest_buf;
 		ts_header_args.dest_offset        = args.dest_offset + nbytes;
 		ts_header_args.dest_buf_sz        = args.dest_buf_sz;
 		ts_header_args.pid                = args.pid;
+		ts_header_args.pcr		  = args.pcr;
 		ts_header_args.continuity_counter = args.continuity_counter;
 		ts_header_args.wrote_pes_header   = wrote_pes_header;
-		ts_header_args.n_stuffing_bytes   = num_stuffing_bytes;
+		ts_header_args.n_stuffing_bytes   = stuff_bytes;
+
+		nbytes += vidtv_pes_write_ts_h(ts_header_args, need_pcr,
+					       &last_pcr);
 
-		nbytes += vidtv_pes_write_ts_h(ts_header_args);
+		need_pcr = false;
 
 		if (!wrote_pes_header) {
 			/* write the PES header only once */
@@ -375,11 +427,11 @@ u32 vidtv_pes_write_into(struct pes_write_args args)
 				       args.dest_offset + nbytes,
 				       args.dest_buf_sz,
 				       args.from,
-				       payload_write_len);
+				       payload_size);
 
-		args.from += payload_write_len;
+		args.from += payload_size;
 
-		remaining_len -= payload_write_len;
+		remaining_len -= payload_size;
 	}
 
 	return nbytes;
diff --git a/drivers/media/test-drivers/vidtv/vidtv_pes.h b/drivers/media/test-drivers/vidtv/vidtv_pes.h
index 84628780d8c6..0ea9e863024d 100644
--- a/drivers/media/test-drivers/vidtv/vidtv_pes.h
+++ b/drivers/media/test-drivers/vidtv/vidtv_pes.h
@@ -125,6 +125,7 @@ struct pes_ts_header_write_args {
 	u8 *continuity_counter;
 	bool wrote_pes_header;
 	u32 n_stuffing_bytes;
+	u64 pcr;
 };
 
 /**
@@ -168,6 +169,7 @@ struct pes_write_args {
 	u64 dts;
 
 	u32 n_pes_h_s_bytes;
+	u64 pcr;
 };
 
 /**