channel_layout: add support for Ambisonic

Signed-off-by: James Almer <jamrial@gmail.com>

3 files changed, 269 insertions, 5 deletions

diff --git a/libavutil/channel_layout.c b/libavutil/channel_layout.c
index 506604cc9f..05ed35c078 100644
--- a/libavutil/channel_layout.c
+++ b/libavutil/channel_layout.c
@@ -34,6 +34,9 @@
 #include "macros.h"
 #include "opt.h"
 
+#define CHAN_IS_AMBI(x) ((x) >= AV_CHAN_AMBISONIC_BASE &&\
+                         (x) <= AV_CHAN_AMBISONIC_END)
+
 struct channel_name {
     const char *name;
     const char *description;
@@ -82,7 +85,10 @@ static const char *get_channel_name(enum AVChannel channel_id)
 
 void av_channel_name_bprint(AVBPrint *bp, enum AVChannel channel_id)
 {
-    if ((unsigned)channel_id < FF_ARRAY_ELEMS(channel_names))
+    if (channel_id >= AV_CHAN_AMBISONIC_BASE &&
+        channel_id <= AV_CHAN_AMBISONIC_END)
+        av_bprintf(bp, "AMBI%d", channel_id - AV_CHAN_AMBISONIC_BASE);
+    else if ((unsigned)channel_id < FF_ARRAY_ELEMS(channel_names))
         av_bprintf(bp, "%s", channel_names[channel_id].name);
     else if (channel_id == AV_CHAN_NONE)
         av_bprintf(bp, "NONE");
@@ -105,7 +111,10 @@ int av_channel_name(char *buf, size_t buf_size, enum AVChannel channel_id)
 
 void av_channel_description_bprint(AVBPrint *bp, enum AVChannel channel_id)
 {
-    if ((unsigned)channel_id < FF_ARRAY_ELEMS(channel_names))
+    if (channel_id >= AV_CHAN_AMBISONIC_BASE &&
+        channel_id <= AV_CHAN_AMBISONIC_END)
+        av_bprintf(bp, "ambisonic ACN %d", channel_id - AV_CHAN_AMBISONIC_BASE);
+    else if ((unsigned)channel_id < FF_ARRAY_ELEMS(channel_names))
         av_bprintf(bp, "%s", channel_names[channel_id].description);
     else
         av_bprintf(bp, "user %d", channel_id);
@@ -129,6 +138,14 @@ enum AVChannel av_channel_from_string(const char *str)
     int i;
     char *endptr = (char *)str;
     enum AVChannel id = AV_CHAN_NONE;
+
+    if (!strncmp(str, "AMBI", 4)) {
+        i = strtol(str + 4, NULL, 0);
+        if (i < 0 || i > AV_CHAN_AMBISONIC_END - AV_CHAN_AMBISONIC_BASE)
+            return AV_CHAN_NONE;
+        return AV_CHAN_AMBISONIC_BASE + i;
+    }
+
     for (i = 0; i < FF_ARRAY_ELEMS(channel_names); i++) {
         if (channel_names[i].name && !strcmp(str, channel_names[i].name))
             return i;
@@ -396,6 +413,65 @@ int av_channel_layout_from_string(AVChannelLayout *channel_layout,
         }
     }
 
+    /* ambisonic */
+    if (!strncmp(str, "ambisonic ", 10)) {
+        const char *p = str + 10;
+        char *endptr;
+        AVChannelLayout extra = {0};
+        int order;
+
+        order = strtol(p, &endptr, 0);
+        if (order < 0 || order + 1  > INT_MAX / (order + 1) ||
+            (*endptr && *endptr != '+'))
+            return AVERROR(EINVAL);
+
+        channel_layout->order       = AV_CHANNEL_ORDER_AMBISONIC;
+        channel_layout->nb_channels = (order + 1) * (order + 1);
+
+        if (*endptr) {
+            int ret = av_channel_layout_from_string(&extra, endptr + 1);
+            if (ret < 0)
+                return ret;
+            if (extra.nb_channels >= INT_MAX - channel_layout->nb_channels) {
+                av_channel_layout_uninit(&extra);
+                return AVERROR(EINVAL);
+            }
+
+            if (extra.order == AV_CHANNEL_ORDER_NATIVE) {
+                channel_layout->u.mask = extra.u.mask;
+            } else {
+                channel_layout->order = AV_CHANNEL_ORDER_CUSTOM;
+                channel_layout->u.map =
+                    av_calloc(channel_layout->nb_channels + extra.nb_channels,
+                              sizeof(*channel_layout->u.map));
+                if (!channel_layout->u.map) {
+                    av_channel_layout_uninit(&extra);
+                    return AVERROR(ENOMEM);
+                }
+
+                for (i = 0; i < channel_layout->nb_channels; i++)
+                    channel_layout->u.map[i].id = AV_CHAN_AMBISONIC_BASE + i;
+                for (i = 0; i < extra.nb_channels; i++) {
+                    enum AVChannel ch = av_channel_layout_channel_from_index(&extra, i);
+                    if (CHAN_IS_AMBI(ch)) {
+                        av_channel_layout_uninit(&extra);
+                        return AVERROR(EINVAL);
+                    }
+                    channel_layout->u.map[channel_layout->nb_channels + i].id = ch;
+                    if (extra.order == AV_CHANNEL_ORDER_CUSTOM &&
+                        extra.u.map[i].name[0])
+                        av_strlcpy(channel_layout->u.map[channel_layout->nb_channels + i].name,
+                                   extra.u.map[i].name,
+                                   sizeof(channel_layout->u.map[channel_layout->nb_channels + i].name));
+                }
+            }
+            channel_layout->nb_channels += extra.nb_channels;
+            av_channel_layout_uninit(&extra);
+        }
+
+        return 0;
+    }
+
     chlist = av_strdup(str);
     if (!chlist)
         return AVERROR(ENOMEM);
@@ -566,6 +642,77 @@ int av_channel_layout_copy(AVChannelLayout *dst, const AVChannelLayout *src)
     return 0;
 }
 
+/**
+ * If the custom layout is n-th order standard-order ambisonic, with optional
+ * extra non-diegetic channels at the end, write its string description in bp.
+ * Return a negative error code on error.
+ */
+static int try_describe_ambisonic(AVBPrint *bp, const AVChannelLayout *channel_layout)
+{
+    int i, highest_ambi, order;
+
+    highest_ambi = -1;
+    if (channel_layout->order == AV_CHANNEL_ORDER_AMBISONIC)
+        highest_ambi = channel_layout->nb_channels - av_popcount64(channel_layout->u.mask) - 1;
+    else {
+        const AVChannelCustom *map = channel_layout->u.map;
+        for (i = 0; i < channel_layout->nb_channels; i++) {
+            int is_ambi = CHAN_IS_AMBI(map[i].id);
+
+            /* ambisonic following non-ambisonic */
+            if (i > 0 && is_ambi && !CHAN_IS_AMBI(map[i - 1].id))
+                return 0;
+
+            /* non-default ordering */
+            if (is_ambi && map[i].id - AV_CHAN_AMBISONIC_BASE != i)
+                return 0;
+
+            if (CHAN_IS_AMBI(map[i].id))
+                highest_ambi = i;
+        }
+    }
+    /* no ambisonic channels*/
+    if (highest_ambi < 0)
+        return 0;
+
+    order = floor(sqrt(highest_ambi));
+    /* incomplete order - some harmonics are missing */
+    if ((order + 1) * (order + 1) != highest_ambi + 1)
+        return 0;
+
+    av_bprintf(bp, "ambisonic %d", order);
+
+    /* extra channels present */
+    if (highest_ambi < channel_layout->nb_channels - 1) {
+        AVChannelLayout extra = { 0 };
+        char buf[128];
+
+        if (channel_layout->order == AV_CHANNEL_ORDER_AMBISONIC) {
+            extra.order       = AV_CHANNEL_ORDER_NATIVE;
+            extra.nb_channels = av_popcount64(channel_layout->u.mask);
+            extra.u.mask      = channel_layout->u.mask;
+        } else {
+            const AVChannelCustom *map = channel_layout->u.map;
+
+            extra.order       = AV_CHANNEL_ORDER_CUSTOM;
+            extra.nb_channels = channel_layout->nb_channels - highest_ambi - 1;
+            extra.u.map       = av_calloc(extra.nb_channels, sizeof(*extra.u.map));
+            if (!extra.u.map)
+                return AVERROR(ENOMEM);
+
+            memcpy(extra.u.map, &map[highest_ambi + 1],
+                   sizeof(*extra.u.map) * extra.nb_channels);
+        }
+
+        av_channel_layout_describe(&extra, buf, sizeof(buf));
+        av_channel_layout_uninit(&extra);
+
+        av_bprintf(bp, "+%s", buf);
+    }
+
+    return 0;
+}
+
 int av_channel_layout_describe_bprint(const AVChannelLayout *channel_layout,
                                       AVBPrint *bp)
 {
@@ -580,6 +727,11 @@ int av_channel_layout_describe_bprint(const AVChannelLayout *channel_layout,
             }
         // fall-through
     case AV_CHANNEL_ORDER_CUSTOM:
+        if (channel_layout->order == AV_CHANNEL_ORDER_CUSTOM) {
+            int res = try_describe_ambisonic(bp, channel_layout);
+            if (res < 0 || bp->len)
+                return res;
+        }
         if (channel_layout->nb_channels)
             av_bprintf(bp, "%d channels (", channel_layout->nb_channels);
         for (i = 0; i < channel_layout->nb_channels; i++) {
@@ -604,6 +756,8 @@ int av_channel_layout_describe_bprint(const AVChannelLayout *channel_layout,
     case AV_CHANNEL_ORDER_UNSPEC:
         av_bprintf(bp, "%d channels", channel_layout->nb_channels);
         return 0;
+    case AV_CHANNEL_ORDER_AMBISONIC:
+        return try_describe_ambisonic(bp, channel_layout);
     default:
         return AVERROR(EINVAL);
     }
@@ -638,6 +792,13 @@ av_channel_layout_channel_from_index(const AVChannelLayout *channel_layout,
     switch (channel_layout->order) {
     case AV_CHANNEL_ORDER_CUSTOM:
         return channel_layout->u.map[idx].id;
+    case AV_CHANNEL_ORDER_AMBISONIC: {
+        int ambi_channels = channel_layout->nb_channels - av_popcount64(channel_layout->u.mask);
+        if (idx < ambi_channels)
+            return AV_CHAN_AMBISONIC_BASE + idx;
+        idx -= ambi_channels;
+        }
+    // fall-through
     case AV_CHANNEL_ORDER_NATIVE:
         for (i = 0; i < 64; i++) {
             if ((1ULL << i) & channel_layout->u.mask && !idx--)
@@ -674,12 +835,20 @@ int av_channel_layout_index_from_channel(const AVChannelLayout *channel_layout,
             if (channel_layout->u.map[i].id == channel)
                 return i;
         return AVERROR(EINVAL);
+    case AV_CHANNEL_ORDER_AMBISONIC:
     case AV_CHANNEL_ORDER_NATIVE: {
         uint64_t mask = channel_layout->u.mask;
+        int ambi_channels = channel_layout->nb_channels - av_popcount64(mask);
+        if (channel_layout->order == AV_CHANNEL_ORDER_AMBISONIC &&
+            channel >= AV_CHAN_AMBISONIC_BASE) {
+            if (channel - AV_CHAN_AMBISONIC_BASE >= ambi_channels)
+                return AVERROR(EINVAL);
+            return channel - AV_CHAN_AMBISONIC_BASE;
+        }
         if ((unsigned)channel > 63 || !(mask & (1ULL << channel)))
             return AVERROR(EINVAL);
         mask &= (1ULL << channel) - 1;
-        return av_popcount64(mask);
+        return av_popcount64(mask) + ambi_channels;
         }
     default:
         return AVERROR(EINVAL);
@@ -711,6 +880,7 @@ int av_channel_layout_index_from_string(const AVChannelLayout *channel_layout,
                 return i;
         }
         // fall-through
+    case AV_CHANNEL_ORDER_AMBISONIC:
     case AV_CHANNEL_ORDER_NATIVE:
         ch = av_channel_from_string(str);
         if (ch == AV_CHAN_NONE)
@@ -737,6 +907,9 @@ int av_channel_layout_check(const AVChannelLayout *channel_layout)
                 return 0;
         }
         return 1;
+    case AV_CHANNEL_ORDER_AMBISONIC:
+        /* If non-diegetic channels are present, ensure they are taken into account */
+        return av_popcount64(channel_layout->u.mask) < channel_layout->nb_channels;
     case AV_CHANNEL_ORDER_UNSPEC:
         return 1;
     default:
@@ -761,7 +934,8 @@ int av_channel_layout_compare(const AVChannelLayout *chl, const AVChannelLayout
         return 0;
 
     /* can compare masks directly */
-    if (chl->order != AV_CHANNEL_ORDER_CUSTOM &&
+    if ((chl->order == AV_CHANNEL_ORDER_NATIVE ||
+         chl->order == AV_CHANNEL_ORDER_AMBISONIC) &&
         chl->order == chl1->order)
         return chl->u.mask != chl1->u.mask;
 
diff --git a/libavutil/channel_layout.h b/libavutil/channel_layout.h
index b1456c245e..4dd6614de9 100644
--- a/libavutil/channel_layout.h
+++ b/libavutil/channel_layout.h
@@ -79,6 +79,23 @@ enum AVChannel {
 
     /** Channel contains data, but its position is unknown. */
     AV_CHAN_UNKNOWN = 0x300,
+
+    /**
+     * Range of channels between AV_CHAN_AMBISONIC_BASE and
+     * AV_CHAN_AMBISONIC_END represent Ambisonic components using the ACN system.
+     *
+     * Given a channel id <i> between AV_CHAN_AMBISONIC_BASE and
+     * AV_CHAN_AMBISONIC_END (inclusive), the ACN index of the channel <n> is
+     * <n> = <i> - AV_CHAN_AMBISONIC_BASE.
+     *
+     * @note these values are only used for AV_CHANNEL_ORDER_CUSTOM channel
+     * orderings, the AV_CHANNEL_ORDER_AMBISONIC ordering orders the channels
+     * implicitly by their position in the stream.
+     */
+    AV_CHAN_AMBISONIC_BASE = 0x400,
+    // leave space for 1024 ids, which correspond to maximum order-32 harmonics,
+    // which should be enough for the foreseeable use cases
+    AV_CHAN_AMBISONIC_END  = 0x7ff,
 };
 
 enum AVChannelOrder {
@@ -100,6 +117,29 @@ enum AVChannelOrder {
      * channels at arbitrary positions.
      */
     AV_CHANNEL_ORDER_CUSTOM,
+    /**
+     * The audio is represented as the decomposition of the sound field into
+     * spherical harmonics. Each channel corresponds to a single expansion
+     * component. Channels are ordered according to ACN (Ambisonic Channel
+     * Number).
+     *
+     * The channel with the index n in the stream contains the spherical
+     * harmonic of degree l and order m given by
+     * @code{.unparsed}
+     *   l   = floor(sqrt(n)),
+     *   m   = n - l * (l + 1).
+     * @endcode
+     *
+     * Conversely given a spherical harmonic of degree l and order m, the
+     * corresponding channel index n is given by
+     * @code{.unparsed}
+     *   n = l * (l + 1) + m.
+     * @endcode
+     *
+     * Normalization is assumed to be SN3D (Schmidt Semi-Normalization)
+     * as defined in AmbiX format $ 2.1.
+     */
+    AV_CHANNEL_ORDER_AMBISONIC,
 };
 
 
@@ -266,7 +306,8 @@ typedef struct AVChannelLayout {
      */
     union {
         /**
-         * This member must be used for AV_CHANNEL_ORDER_NATIVE.
+         * This member must be used for AV_CHANNEL_ORDER_NATIVE, and may be used
+         * for AV_CHANNEL_ORDER_AMBISONIC to signal non-diegetic channels.
          * It is a bitmask, where the position of each set bit means that the
          * AVChannel with the corresponding value is present.
          *
@@ -288,6 +329,11 @@ typedef struct AVChannelLayout {
          * I.e. when map[i].id is equal to AV_CHAN_FOO, then AV_CH_FOO is the
          * i-th channel in the audio data.
          *
+         * When map[i].id is in the range between AV_CHAN_AMBISONIC_BASE and
+         * AV_CHAN_AMBISONIC_END (inclusive), the channel contains an ambisonic
+         * component with ACN index (as defined above)
+         * n = map[i].id - AV_CHAN_AMBISONIC_BASE.
+         *
          * map[i].name may be filled with a 0-terminated string, in which case
          * it will be used for the purpose of identifying the channel with the
          * convenience functions below. Otherise it must be zeroed.
@@ -333,6 +379,8 @@ typedef struct AVChannelLayout {
 #define AV_CHANNEL_LAYOUT_HEXADECAGONAL     AV_CHANNEL_LAYOUT_MASK(16, AV_CH_LAYOUT_HEXADECAGONAL)
 #define AV_CHANNEL_LAYOUT_STEREO_DOWNMIX    AV_CHANNEL_LAYOUT_MASK(2,  AV_CH_LAYOUT_STEREO_DOWNMIX)
 #define AV_CHANNEL_LAYOUT_22POINT2          AV_CHANNEL_LAYOUT_MASK(24, AV_CH_LAYOUT_22POINT2)
+#define AV_CHANNEL_LAYOUT_AMBISONIC_FIRST_ORDER \
+    { .order = AV_CHANNEL_ORDER_AMBISONIC, .nb_channels = 4, .u = { .mask = 0 }}
 
 struct AVBPrint;
 
@@ -532,6 +580,8 @@ int av_channel_layout_from_mask(AVChannelLayout *channel_layout, uint64_t mask);
  *  - a decimal or hexadecimal value of a native channel layout (eg. "4" or "0x4")
  *  - the number of channels with default layout (eg. "4c")
  *  - the number of unordered channels (eg. "4C" or "4 channels")
+ *  - the ambisonic order followed by optional non-diegetic channels (eg.
+ *    "ambisonic 2+stereo")
  *
  * @param channel_layout input channel layout
  * @param str string describing the channel layout
diff --git a/libavutil/tests/channel_layout.c b/libavutil/tests/channel_layout.c
index 83d549a42c..4830c4c730 100644
--- a/libavutil/tests/channel_layout.c
+++ b/libavutil/tests/channel_layout.c
@@ -79,6 +79,10 @@ int main(void)
     printf("With AV_CHAN_FRONT_RIGHT: %26s\n", bp.str);
     CHANNEL_NAME(63);
     printf("With 63: %43s\n", bp.str);
+    CHANNEL_NAME(AV_CHAN_AMBISONIC_BASE);
+    printf("With AV_CHAN_AMBISONIC_BASE: %23s\n", bp.str);
+    CHANNEL_NAME(AV_CHAN_AMBISONIC_END);
+    printf("With AV_CHAN_AMBISONIC_END: %24s\n", bp.str);
 
     printf("Testing av_channel_description\n");
     CHANNEL_DESCRIPTION(AV_CHAN_FRONT_LEFT);
@@ -87,11 +91,17 @@ int main(void)
     printf("With AV_CHAN_FRONT_RIGHT: %26s\n", bp.str);
     CHANNEL_DESCRIPTION(63);
     printf("With 63: %43s\n", bp.str);
+    CHANNEL_DESCRIPTION(AV_CHAN_AMBISONIC_BASE);
+    printf("With AV_CHAN_AMBISONIC_BASE: %23s\n", bp.str);
+    CHANNEL_DESCRIPTION(AV_CHAN_AMBISONIC_END);
+    printf("With AV_CHAN_AMBISONIC_END: %24s\n", bp.str);
 
     printf("\nTesting av_channel_from_string\n");
     printf("With \"FL\": %41d\n", av_channel_from_string("FL"));
     printf("With \"FR\": %41d\n", av_channel_from_string("FR"));
     printf("With \"USR63\": %38d\n", av_channel_from_string("USR63"));
+    printf("With \"AMBI0\": %38d\n", av_channel_from_string("AMBI0"));
+    printf("With \"AMBI1023\": %35d\n", av_channel_from_string("AMBI1023"));
 
     printf("\n==Native layouts==\n");
 
@@ -194,6 +204,8 @@ int main(void)
     printf("With \"FL+FR+FC+BL+BR+LFE\": %34s\n", bp.str);
     CHANNEL_LAYOUT_FROM_STRING("2 channels (FR+FL)");
     printf("With \"2 channels (FR+FL)\": %34s\n", bp.str);
+    CHANNEL_LAYOUT_FROM_STRING("ambisonic 1+FR+FL");
+    printf("With \"ambisonic 1+FR+FL\": %35s\n", bp.str);
     CHANNEL_LAYOUT_FROM_STRING("FL@Foo+FR@Bar");
     printf("With \"FL@Foo+FR@Bar\": %39s\n", bp.str);
     CHANNEL_LAYOUT_FROM_STRING("FR+FL@Foo+USR63@Foo");
@@ -259,6 +271,34 @@ int main(void)
     CHANNEL_LAYOUT_CHANNEL_FROM_INDEX(3);
     printf("On \"FR+FL@Foo+USR63@Foo\" layout with 3: %21d\n", ret);
 
+    printf("\n==Ambisonic layouts==\n");
+
+    printf("\nTesting av_channel_layout_from_string\n");
+    CHANNEL_LAYOUT_FROM_STRING("ambisonic 1");
+    printf("With \"ambisonic 1\": %41s\n", bp.str);
+    CHANNEL_LAYOUT_FROM_STRING("ambisonic 2+stereo");
+    printf("With \"ambisonic 2+stereo\": %34s\n", bp.str);
+
+    printf("\nTesting av_channel_layout_index_from_channel\n");
+    CHANNEL_LAYOUT_INDEX_FROM_CHANNEL(AV_CHAN_AMBISONIC_BASE);
+    printf("On \"ambisonic 2+stereo\" layout with AV_CHAN_AMBISONIC_BASE: %d\n", ret);
+    CHANNEL_LAYOUT_INDEX_FROM_CHANNEL(AV_CHAN_FRONT_LEFT);
+    printf("On \"ambisonic 2+stereo\" layout with AV_CHAN_FRONT_LEFT: %5d\n", ret);
+    CHANNEL_LAYOUT_INDEX_FROM_CHANNEL(AV_CHAN_FRONT_RIGHT);
+    printf("On \"ambisonic 2+stereo\" layout with AV_CHAN_FRONT_RIGHT: %4d\n", ret);
+    CHANNEL_LAYOUT_INDEX_FROM_CHANNEL(AV_CHAN_BACK_CENTER);
+    printf("On \"ambisonic 2+stereo\" layout with AV_CHAN_BACK_CENTER: %4d\n", ret);
+
+    printf("\nTesting av_channel_layout_channel_from_index\n");
+    CHANNEL_LAYOUT_CHANNEL_FROM_INDEX(0);
+    printf("On \"ambisonic 2+stereo\" layout with 0: %22d\n", ret);
+    CHANNEL_LAYOUT_CHANNEL_FROM_INDEX(9);
+    printf("On \"ambisonic 2+stereo\" layout with 9: %22d\n", ret);
+    CHANNEL_LAYOUT_CHANNEL_FROM_INDEX(10);
+    printf("On \"ambisonic 2+stereo\" layout with 10: %21d\n", ret);
+    CHANNEL_LAYOUT_CHANNEL_FROM_INDEX(11);
+    printf("On \"ambisonic 2+stereo\" layout with 11: %21d\n", ret);
+
     av_channel_layout_uninit(&layout);
     av_bprint_finalize(&bp, NULL);