Splitting up DRED encoding

4 files changed, 31 insertions, 32 deletions

diff --git a/silk/dred_encoder.c b/silk/dred_encoder.c
index a5d6b0d4..9e9077f8 100644
--- a/silk/dred_encoder.c
+++ b/silk/dred_encoder.c
@@ -49,21 +49,11 @@ void init_dred_encoder(DREDEnc* enc)
     enc->rdovae_enc = DRED_rdovae_create_encoder();
 }
 
-void dred_encode_silk_frame(DREDEnc *enc, const opus_int16 *silk_frame)
+void dred_process_silk_frame(DREDEnc *enc, const opus_int16 *silk_frame)
 {
-    const opus_uint16 *dead_zone       = DRED_rdovae_get_dead_zone_pointer();
-    const opus_uint16 *p0              = DRED_rdovae_get_p0_pointer();
-    const opus_uint16 *quant_scales    = DRED_rdovae_get_quant_scales_pointer();
-    const opus_uint16 *r               = DRED_rdovae_get_r_pointer();
     float feature_buffer[2 * 36];
 
     float input_buffer[2*DRED_NUM_FEATURES] = {0};
-    ec_enc ec_encoder;
-
-    int q_level;
-    int i;
-    int offset;
-
     /* delay signal by 79 samples */
     memmove(enc->input_buffer, enc->input_buffer + DRED_DFRAME_SIZE, DRED_SILK_ENCODER_DELAY * sizeof(*enc->input_buffer));
     memcpy(enc->input_buffer + DRED_SILK_ENCODER_DELAY, silk_frame, DRED_DFRAME_SIZE * sizeof(*silk_frame));
@@ -82,9 +72,22 @@ void dred_encode_silk_frame(DREDEnc *enc, const opus_int16 *silk_frame)
     /* run RDOVAE encoder */
     DRED_rdovae_encode_dframe(enc->rdovae_enc, enc->latents_buffer, enc->state_buffer, input_buffer);
     enc->latents_buffer_fill = IMIN(enc->latents_buffer_fill+1, DRED_NUM_REDUNDANCY_FRAMES);
+}
+
+int dred_encode_silk_frame(DREDEnc *enc, unsigned char *buf, int max_frames, int max_bytes) {
+    const opus_uint16 *dead_zone       = DRED_rdovae_get_dead_zone_pointer();
+    const opus_uint16 *p0              = DRED_rdovae_get_p0_pointer();
+    const opus_uint16 *quant_scales    = DRED_rdovae_get_quant_scales_pointer();
+    const opus_uint16 *r               = DRED_rdovae_get_r_pointer();
+    ec_enc ec_encoder;
+
+    int q_level;
+    int i;
+    int offset;
+    int ec_buffer_fill;
 
     /* entropy coding of state and latents */
-    ec_enc_init(&ec_encoder, enc->ec_buffer, DRED_MAX_DATA_SIZE);
+    ec_enc_init(&ec_encoder, buf, max_bytes);
     dred_encode_state(&ec_encoder, enc->state_buffer);
 
     for (i = 0; i < enc->latents_buffer_fill-1; i += 2)
@@ -102,20 +105,8 @@ void dred_encode_silk_frame(DREDEnc *enc, const opus_int16 *silk_frame)
         );
     }
 
-    enc->ec_buffer_fill = (ec_tell(&ec_encoder)+7)/8;
-    ec_enc_shrink(&ec_encoder, enc->ec_buffer_fill);
+    ec_buffer_fill = (ec_tell(&ec_encoder)+7)/8;
+    ec_enc_shrink(&ec_encoder, ec_buffer_fill);
     ec_enc_done(&ec_encoder);
-
-#if 0
-    printf("packet size: %d\n", enc->ec_buffer_fill*8);
-
-    static FILE *fbs = NULL;
-    if (fbs == NULL)
-    {
-        fbs = fopen("dred_bitstream.bin", "wb");
-    }
-    fwrite(&enc->ec_buffer_fill, sizeof(enc->ec_buffer_fill), 1, fbs);
-    fwrite(ec_get_buffer(&ec_encoder), 1, enc->ec_buffer_fill, fbs);
-#endif
-
+    return ec_buffer_fill;
 }
diff --git a/silk/dred_encoder.h b/silk/dred_encoder.h
index 9bb96b96..a32cfa97 100644
--- a/silk/dred_encoder.h
+++ b/silk/dred_encoder.h
@@ -40,8 +40,6 @@ typedef struct {
     float latents_buffer[DRED_MAX_FRAMES * DRED_LATENT_DIM];
     int latents_buffer_fill;
     float state_buffer[24];
-    unsigned char ec_buffer[DRED_MAX_DATA_SIZE];
-    int ec_buffer_fill;
     LPCNetEncState *lpcnet_enc_state;
     RDOVAEEnc *rdovae_enc;
 } DREDEnc;
@@ -51,6 +49,8 @@ void init_dred_encoder(DREDEnc* enc);
 
 void dred_deinit_encoder(DREDEnc *enc);
 
-void dred_encode_silk_frame(DREDEnc *enc, const opus_int16 *silk_frame);
+void dred_process_silk_frame(DREDEnc *enc, const opus_int16 *silk_frame);
+
+int dred_encode_silk_frame(DREDEnc *enc, unsigned char *buf, int max_frames, int max_bytes);
 
 #endif
diff --git a/silk/enc_API.c b/silk/enc_API.c
index 05ce9cb8..45d67877 100644
--- a/silk/enc_API.c
+++ b/silk/enc_API.c
@@ -470,7 +470,7 @@ opus_int silk_Encode(                                   /* O    Returns error co
 
 #ifdef ENABLE_NEURAL_FEC
             /* DRED Encoder */
-            dred_encode_silk_frame( &psEnc->state_Fxx[ 0 ].sCmn.dred_encoder, &psEnc->state_Fxx[ 0 ].sCmn.inputBuf[0] );
+            dred_process_silk_frame( &psEnc->state_Fxx[ 0 ].sCmn.dred_encoder, &psEnc->state_Fxx[ 0 ].sCmn.inputBuf[0] );
 #endif
 
             /* Encode */
diff --git a/src/opus_encoder.c b/src/opus_encoder.c
index 4354a249..47a4d6eb 100644
--- a/src/opus_encoder.c
+++ b/src/opus_encoder.c
@@ -2181,8 +2181,16 @@ opus_int32 opus_encode_native(OpusEncoder *st, const opus_val16 *pcm, int frame_
     ret += 1+redundancy_bytes;
 #ifdef ENABLE_NEURAL_FEC
     if (1) {
+       opus_extension_data extension;
+       unsigned char buf[DRED_MAX_DATA_SIZE];
+       int dred_bytes;
        DREDEnc *dred = &((silk_encoder*)silk_enc)->state_Fxx[0].sCmn.dred_encoder;
-       opus_extension_data extension = {127, 0, dred->ec_buffer, dred->ec_buffer_fill};
+       dred_bytes = IMIN(DRED_MAX_DATA_SIZE, max_data_bytes-ret-2);
+       dred_bytes = dred_encode_silk_frame(dred, buf, DRED_NUM_REDUNDANCY_FRAMES/2, dred_bytes);
+       extension.id = 127;
+       extension.frame = 0;
+       extension.data = buf;
+       extension.len = dred_bytes;
        ret = opus_packet_pad_impl(data, ret, max_data_bytes, !st->use_vbr, &extension, 1);
        if (ret < 0)
        {