diff --git a/libavcodec/cbs_h2645.c b/libavcodec/cbs_h2645.c
index 6005d46e0d..36212d1da6 100644
--- a/libavcodec/cbs_h2645.c
+++ b/libavcodec/cbs_h2645.c
@@ -661,38 +661,127 @@ static int cbs_h2645_split_fragment(CodedBitstreamContext *ctx,
     return 0;
 }
 
-#define cbs_h2645_replace_ps(h26n, ps_name, ps_var, id_element) \
-static int cbs_h26 ## h26n ## _replace_ ## ps_var(CodedBitstreamContext *ctx, \
-                                                  CodedBitstreamUnit *unit)  \
-{ \
-    CodedBitstreamH26 ## h26n ## Context *priv = ctx->priv_data; \
-    H26 ## h26n ## Raw ## ps_name *ps_var = unit->content; \
-    unsigned int id = ps_var->id_element; \
-    int err; \
-    if (id >= FF_ARRAY_ELEMS(priv->ps_var)) { \
-        av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid " #ps_name \
-               " id : %d.\n", id); \
-        return AVERROR_INVALIDDATA; \
-    } \
-    err = ff_cbs_make_unit_refcounted(ctx, unit); \
-    if (err < 0) \
-        return err; \
-    if (priv->ps_var[id] == priv->active_ ## ps_var) \
-        priv->active_ ## ps_var = NULL ; \
-    av_buffer_unref(&priv->ps_var ## _ref[id]); \
-    av_assert0(unit->content_ref); \
-    priv->ps_var ## _ref[id] = av_buffer_ref(unit->content_ref); \
-    if (!priv->ps_var ## _ref[id]) \
-        return AVERROR(ENOMEM); \
-    priv->ps_var[id] = (H26 ## h26n ## Raw ## ps_name *)priv->ps_var ## _ref[id]->data; \
-    return 0; \
-}
+static int cbs_h2645_replace_ps(CodedBitstreamContext *ctx,
+                                CodedBitstreamUnit *unit)
+{
+    typedef struct {
+        int nal_unit_type;
+        int max_count;
+        const char *name;
+        size_t id_offset;
+        size_t ref_array_offset;
+        size_t ptr_array_offset;
+        size_t active_offset;
+    } PSType;
+
+    static const PSType h264_ps_types[] = {
+        {
+            H264_NAL_SPS,
+            H264_MAX_SPS_COUNT,
+            "SPS",
+            offsetof(H264RawSPS, seq_parameter_set_id),
+            offsetof(CodedBitstreamH264Context, sps_ref),
+            offsetof(CodedBitstreamH264Context, sps),
+            offsetof(CodedBitstreamH264Context, active_sps),
+        },
+        {
+            H264_NAL_PPS,
+            H264_MAX_PPS_COUNT,
+            "PPS",
+            offsetof(H264RawPPS, pic_parameter_set_id),
+            offsetof(CodedBitstreamH264Context, pps_ref),
+            offsetof(CodedBitstreamH264Context, pps),
+            offsetof(CodedBitstreamH264Context, active_pps),
+        },
+    };
+
+    static const PSType h265_ps_types[] = {
+        {
+            HEVC_NAL_VPS,
+            HEVC_MAX_VPS_COUNT,
+            "VPS",
+            offsetof(H265RawVPS, vps_video_parameter_set_id),
+            offsetof(CodedBitstreamH265Context, vps_ref),
+            offsetof(CodedBitstreamH265Context, vps),
+            offsetof(CodedBitstreamH265Context, active_vps),
+        },
+        {
+            HEVC_NAL_SPS,
+            HEVC_MAX_SPS_COUNT,
+            "SPS",
+            offsetof(H265RawSPS, sps_seq_parameter_set_id),
+            offsetof(CodedBitstreamH265Context, sps_ref),
+            offsetof(CodedBitstreamH265Context, sps),
+            offsetof(CodedBitstreamH265Context, active_sps),
+        },
+        {
+            HEVC_NAL_PPS,
+            HEVC_MAX_PPS_COUNT,
+            "PPS",
+            offsetof(H265RawPPS, pps_pic_parameter_set_id),
+            offsetof(CodedBitstreamH265Context, pps_ref),
+            offsetof(CodedBitstreamH265Context, pps),
+            offsetof(CodedBitstreamH265Context, active_pps),
+        },
+    };
 
-cbs_h2645_replace_ps(4, SPS, sps, seq_parameter_set_id)
-cbs_h2645_replace_ps(4, PPS, pps, pic_parameter_set_id)
-cbs_h2645_replace_ps(5, VPS, vps, vps_video_parameter_set_id)
-cbs_h2645_replace_ps(5, SPS, sps, sps_seq_parameter_set_id)
-cbs_h2645_replace_ps(5, PPS, pps, pps_pic_parameter_set_id)
+    const PSType *ps_type;
+    AVBufferRef **ref_array;
+    void **ptr_array;
+    void **active;
+    int err, id, i, nb_ps_types;
+
+    switch (ctx->codec->codec_id) {
+    case AV_CODEC_ID_H264:
+        ps_type     = h264_ps_types;
+        nb_ps_types = FF_ARRAY_ELEMS(h264_ps_types);
+        break;
+    case AV_CODEC_ID_H265:
+        ps_type     = h265_ps_types;
+        nb_ps_types = FF_ARRAY_ELEMS(h265_ps_types);
+        break;
+    default:
+        av_assert0(0);
+    }
+
+    for (i = 0; i < nb_ps_types; i++) {
+        if (ps_type->nal_unit_type == unit->type)
+            break;
+        ++ps_type;
+    }
+    av_assert0(i < nb_ps_types);
+
+    id = *((uint8_t*)unit->content + ps_type->id_offset);
+
+    if (id >= ps_type->max_count) {
+        av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid %s id: %d.\n",
+               ps_type->name, id);
+        return AVERROR_INVALIDDATA;
+    }
+
+    err = ff_cbs_make_unit_refcounted(ctx, unit);
+    if (err < 0)
+        return err;
+
+    ref_array =
+         (AVBufferRef**)((uint8_t*)ctx->priv_data + ps_type->ref_array_offset);
+    ptr_array = (void**)((uint8_t*)ctx->priv_data + ps_type->ptr_array_offset);
+    active    = (void**)((uint8_t*)ctx->priv_data + ps_type->active_offset);
+
+    if (ptr_array[id] == *active) {
+        // The old active parameter set is being overwritten, so it can't
+        // be active after this point.
+        *active = NULL;
+    }
+    av_buffer_unref(&ref_array[id]);
+
+    ref_array[id] = av_buffer_ref(unit->content_ref);
+    if (!ref_array[id])
+        return AVERROR(ENOMEM);
+    ptr_array[id] = ref_array[id]->data;
+
+    return 0;
+}
 
 static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
                                   CodedBitstreamUnit *unit)
@@ -717,7 +806,7 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h264_replace_sps(ctx, unit);
+            err = cbs_h2645_replace_ps(ctx, unit);
             if (err < 0)
                 return err;
         }
@@ -739,7 +828,7 @@ static int cbs_h264_read_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h264_replace_pps(ctx, unit);
+            err = cbs_h2645_replace_ps(ctx, unit);
             if (err < 0)
                 return err;
         }
@@ -836,7 +925,7 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h265_replace_vps(ctx, unit);
+            err = cbs_h2645_replace_ps(ctx, unit);
             if (err < 0)
                 return err;
         }
@@ -849,7 +938,7 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h265_replace_sps(ctx, unit);
+            err = cbs_h2645_replace_ps(ctx, unit);
             if (err < 0)
                 return err;
         }
@@ -863,7 +952,7 @@ static int cbs_h265_read_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h265_replace_pps(ctx, unit);
+            err = cbs_h2645_replace_ps(ctx, unit);
             if (err < 0)
                 return err;
         }
@@ -1007,7 +1096,7 @@ static int cbs_h264_write_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h264_replace_sps(ctx, unit);
+            err = cbs_h2645_replace_ps(ctx, unit);
             if (err < 0)
                 return err;
         }
@@ -1031,7 +1120,7 @@ static int cbs_h264_write_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h264_replace_pps(ctx, unit);
+            err = cbs_h2645_replace_ps(ctx, unit);
             if (err < 0)
                 return err;
         }
@@ -1124,7 +1213,7 @@ static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h265_replace_vps(ctx, unit);
+            err = cbs_h2645_replace_ps(ctx, unit);
             if (err < 0)
                 return err;
         }
@@ -1138,7 +1227,7 @@ static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h265_replace_sps(ctx, unit);
+            err = cbs_h2645_replace_ps(ctx, unit);
             if (err < 0)
                 return err;
         }
@@ -1152,7 +1241,7 @@ static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
             if (err < 0)
                 return err;
 
-            err = cbs_h265_replace_pps(ctx, unit);
+            err = cbs_h2645_replace_ps(ctx, unit);
             if (err < 0)
                 return err;
         }
