Add support for -D in qsb and preamble in QShaderBaker

Allow inserting a user provided string. glslang takes care of
heavy lifting. For qsb this means compile time #define statements
can now be generated by passing -DSOMETHING and -DSOMETHING=123 on
the command line.

Change-Id: Iddf276c3fa91e21780530aa2f54719325f404ba6
Reviewed-by: Laszlo Agocs <laszlo.agocs@qt.io>

8 files changed, 154 insertions, 5 deletions

diff --git a/src/shadertools/qshaderbaker.cpp b/src/shadertools/qshaderbaker.cpp
index fbb0a2b..982a2fd 100644
--- a/src/shadertools/qshaderbaker.cpp
+++ b/src/shadertools/qshaderbaker.cpp
@@ -150,6 +150,7 @@ struct QShaderBakerPrivate
     QRhiShader::ShaderStage stage;
     QVector<QShaderBaker::GeneratedShader> reqVersions;
     QVector<QRhiShaderKey::ShaderVariant> variants;
+    QByteArray preamble;
     QSpirvCompiler compiler;
     QString errorMessage;
 };
@@ -281,7 +282,7 @@ void QShaderBaker::setGeneratedShaders(const QVector<GeneratedShader> &v)
 }
 
 /*!
-    Specifies which shader variants are genetated. Each shader version can have
+    Specifies which shader variants are generated. Each shader version can have
     multiple variants in the resulting QRhiShader.
 
     In most cases \a v contains a single entry, QRhiShaderKey::StandardShader.
@@ -295,6 +296,23 @@ void QShaderBaker::setGeneratedShaderVariants(const QVector<QRhiShaderKey::Shade
 }
 
 /*!
+    Specifies a custom \a preamble that is processed before the normal shader
+    code.
+
+    This is more than just prepending to the source string: the validity of the
+    GLSL version directive, which is required to be placed before everything
+    else, is not affected. Line numbers in the reported error messages also
+    remain unchanged, ignoring the contents given in the \a preamble.
+
+    One use case for preambles is to transparently insert dynamically generated
+    \c{#define} statements.
+ */
+void QShaderBaker::setPreamble(const QByteArray &preamble)
+{
+    d->preamble = preamble;
+}
+
+/*!
     Runs the compilation and translation process.
 
     \return a QRhiShader instance. To check if the process was successful,
@@ -319,6 +337,7 @@ QRhiShader QShaderBaker::bake()
 
     d->compiler.setSourceString(d->source, d->stage, d->sourceFileName);
     d->compiler.setFlags(0);
+    d->compiler.setPreamble(d->preamble);
     QByteArray spirv = d->compiler.compileToSpirv();
     if (spirv.isEmpty()) {
         d->errorMessage = d->compiler.errorMessage();
diff --git a/src/shadertools/qshaderbaker.h b/src/shadertools/qshaderbaker.h
index de6159b..68f0672 100644
--- a/src/shadertools/qshaderbaker.h
+++ b/src/shadertools/qshaderbaker.h
@@ -64,6 +64,8 @@ public:
     void setGeneratedShaders(const QVector<GeneratedShader> &v);
     void setGeneratedShaderVariants(const QVector<QRhiShaderKey::ShaderVariant> &v);
 
+    void setPreamble(const QByteArray &preamble);
+
     QRhiShader bake();
 
     QString errorMessage() const;
diff --git a/src/shadertools/qspirvcompiler.cpp b/src/shadertools/qspirvcompiler.cpp
index 22547e9..d2f5703 100644
--- a/src/shadertools/qspirvcompiler.cpp
+++ b/src/shadertools/qspirvcompiler.cpp
@@ -162,6 +162,7 @@ struct QSpirvCompilerPrivate
     QByteArray batchableSource;
     EShLanguage stage = EShLangVertex;
     QSpirvCompiler::Flags flags = 0;
+    QByteArray preamble;
     QByteArray spirv;
     QString log;
 };
@@ -261,6 +262,11 @@ bool QSpirvCompilerPrivate::compile()
     const char *srcStr = actualSource->constData();
     const int size = actualSource->size();
     shader.setStringsWithLengthsAndNames(&srcStr, &size, &fnStr, 1);
+    if (!preamble.isEmpty()) {
+        // Line numbers in errors and #version are not affected by having a
+        // preamble, which is just what we need.
+        shader.setPreamble(preamble.constData());
+    }
 
     shader.setEnvInput(glslang::EShSourceGlsl, stage, glslang::EShClientVulkan, 100);
     shader.setEnvClient(glslang::EShClientVulkan, glslang::EShTargetVulkan_1_0);
@@ -374,6 +380,11 @@ void QSpirvCompiler::setFlags(Flags flags)
     d->flags = flags;
 }
 
+void QSpirvCompiler::setPreamble(const QByteArray &preamble)
+{
+    d->preamble = preamble;
+}
+
 QByteArray QSpirvCompiler::compileToSpirv()
 {
     if (d->stage == EShLangVertex && d->flags.testFlag(RewriteToMakeBatchableForSG) && d->batchableSource.isEmpty())
diff --git a/src/shadertools/qspirvcompiler_p.h b/src/shadertools/qspirvcompiler_p.h
index 2af56b2..8637795 100644
--- a/src/shadertools/qspirvcompiler_p.h
+++ b/src/shadertools/qspirvcompiler_p.h
@@ -73,6 +73,7 @@ public:
     void setSourceDevice(QIODevice *device, QRhiShader::ShaderStage stage, const QString &fileName = QString());
     void setSourceString(const QByteArray &sourceString, QRhiShader::ShaderStage stage, const QString &fileName = QString());
     void setFlags(Flags flags);
+    void setPreamble(const QByteArray &preamble);
 
     QByteArray compileToSpirv();
     QString errorMessage() const;
diff --git a/tests/auto/qshaderbaker/data/defines.frag b/tests/auto/qshaderbaker/data/defines.frag
new file mode 100644
index 0000000..8bdd794
--- /dev/null
+++ b/tests/auto/qshaderbaker/data/defines.frag
@@ -0,0 +1,28 @@
+#version 440
+
+layout(location = 0) in vec3 v_color;
+layout(location = 0) out vec4 fragColor;
+
+layout(std140, binding = 0) uniform buf {
+    mat4 mvp;
+#if OPACITY_SIZE == 1
+    float opacity;
+#elif OPACITY_SIZE == 2
+    vec2 opacity;
+#elif OPACITY_SIZE == 3
+    vec3 opacity;
+#else
+    vec4 opacity;
+#endif
+} ubuf;
+
+#ifdef DO_NOT_BREAK
+void main()
+#endif
+{
+#if OPACITY_SIZE == 1
+    fragColor = vec4(v_color * ubuf.opacity, ubuf.opacity);
+#else
+    fragColor = vec4(v_color * ubuf.opacity.r, ubuf.opacity.r);
+#endif
+}
diff --git a/tests/auto/qshaderbaker/tst_qshaderbaker.cpp b/tests/auto/qshaderbaker/tst_qshaderbaker.cpp
index 36085e5..e6602d8 100644
--- a/tests/auto/qshaderbaker/tst_qshaderbaker.cpp
+++ b/tests/auto/qshaderbaker/tst_qshaderbaker.cpp
@@ -53,6 +53,7 @@ private slots:
     void compileError();
     void translateError();
     void genVariants();
+    void defines();
 };
 
 void tst_QShaderBaker::initTestCase()
@@ -373,5 +374,56 @@ void tst_QShaderBaker::genVariants()
     QCOMPARE(batchableGlslVariantCount, 3);
 }
 
+void tst_QShaderBaker::defines()
+{
+    QShaderBaker baker;
+    baker.setSourceFileName(QLatin1String(":/data/defines.frag"));
+    baker.setGeneratedShaderVariants({ QRhiShaderKey::StandardShader });
+    baker.setGeneratedShaders({ { QRhiShaderKey::SpirvShader, QRhiShaderVersion(100) } });
+    QRhiShader s = baker.bake();
+    QVERIFY(!s.isValid());
+    QVERIFY(!baker.errorMessage().isEmpty());
+    qDebug() << baker.errorMessage();
+
+    QByteArray preamble;
+    preamble = QByteArrayLiteral("#define DO_NOT_BREAK\n");
+    baker.setPreamble(preamble);
+    s = baker.bake();
+    QVERIFY(s.isValid());
+    QVERIFY(baker.errorMessage().isEmpty());
+
+    QRhiShaderDescription desc = s.description();
+    QCOMPARE(desc.uniformBlocks().count(), 1);
+    QRhiShaderDescription::UniformBlock blk = desc.uniformBlocks().first();
+    QCOMPARE(blk.members.count(), 2);
+    bool opacity_ok = false;
+    for (int i = 0; i < blk.members.count(); ++i) {
+        const QRhiShaderDescription::BlockVariable v = blk.members[i];
+        if (v.name == QLatin1String("opacity")) {
+            opacity_ok = v.type == QRhiShaderDescription::Vec4;
+            break;
+        }
+    }
+    QVERIFY(opacity_ok);
+
+    preamble += QByteArrayLiteral("#define OPACITY_SIZE 1\n");
+    baker.setPreamble(preamble);
+    s = baker.bake();
+    QVERIFY(s.isValid());
+    QVERIFY(baker.errorMessage().isEmpty());
+
+    desc = s.description();
+    blk = desc.uniformBlocks().first();
+    opacity_ok = false;
+    for (int i = 0; i < blk.members.count(); ++i) {
+        const QRhiShaderDescription::BlockVariable v = blk.members[i];
+        if (v.name == QLatin1String("opacity")) {
+            opacity_ok = v.type == QRhiShaderDescription::Float;
+            break;
+        }
+    }
+    QVERIFY(opacity_ok);
+}
+
 #include <tst_qshaderbaker.moc>
 QTEST_MAIN(tst_QShaderBaker)
diff --git a/tests/playground/preamble.frag b/tests/playground/preamble.frag
new file mode 100644
index 0000000..f7b30ef
--- /dev/null
+++ b/tests/playground/preamble.frag
@@ -0,0 +1,17 @@
+#version 440
+
+layout(location = 0) in vec2 v_texcoord;
+layout(location = 0) out vec4 fragColor;
+
+// must be compiled with -DMAKE_IT_WORK -DYES_REALLY=99
+
+#ifdef MAKE_IT_WORK
+#if YES_REALLY > 98
+layout(binding = 1) uniform sampler2D tex;
+#endif
+#endif
+
+void main()
+{
+    fragColor = texture(tex, v_texcoord);
+}
diff --git a/tools/qsb/qsb.cpp b/tools/qsb/qsb.cpp
index 66435c5..ea143f9 100644
--- a/tools/qsb/qsb.cpp
+++ b/tools/qsb/qsb.cpp
@@ -238,25 +238,27 @@ int main(int argc, char **argv)
     cmdLineParser.addOption(batchableOption);
     QCommandLineOption glslOption({ "g", "glsl" },
                                   QObject::tr("Comma separated list of GLSL versions to generate. (for example, \"100 es,120,330\")"),
-                                  QObject::tr("glsl"));
+                                  QObject::tr("versions"));
     cmdLineParser.addOption(glslOption);
     QCommandLineOption hlslOption({ "l", "hlsl" },
                                   QObject::tr("Comma separated list of HLSL (Shader Model) versions to generate. F.ex. 50 is 5.0, 51 is 5.1."),
-                                  QObject::tr("hlsl"));
+                                  QObject::tr("versions"));
     cmdLineParser.addOption(hlslOption);
     QCommandLineOption mslOption({ "m", "msl" },
                                   QObject::tr("Comma separated list of Metal Shading Language versions to generate. F.ex. 12 is 1.2, 20 is 2.0."),
-                                  QObject::tr("msl"));
+                                  QObject::tr("versions"));
     cmdLineParser.addOption(mslOption);
     QCommandLineOption outputOption({ "o", "output" },
                                      QObject::tr("Output file for the baked shader pack."),
-                                     QObject::tr("output"));
+                                     QObject::tr("filename"));
     cmdLineParser.addOption(outputOption);
     QCommandLineOption fxcOption({ "c", "fxc" }, QObject::tr("In combination with --hlsl invokes fxc to store DXBC instead of HLSL."));
     cmdLineParser.addOption(fxcOption);
     QCommandLineOption mtllibOption({ "t", "metallib" },
                                     QObject::tr("In combination with --msl builds a Metal library with xcrun metal(lib) and stores that instead of the source."));
     cmdLineParser.addOption(mtllibOption);
+    QCommandLineOption defineOption({ "D", "define" }, QObject::tr("Define macro"), QObject::tr("name[=value]"));
+    cmdLineParser.addOption(defineOption);
     QCommandLineOption dumpOption({ "d", "dump" }, QObject::tr("Switches to dump mode. Input file is expected to be a baked shader pack."));
     cmdLineParser.addOption(dumpOption);
 
@@ -340,6 +342,23 @@ int main(int argc, char **argv)
 
         baker.setGeneratedShaders(genShaders);
 
+        if (cmdLineParser.isSet(defineOption)) {
+            QByteArray preamble;
+            const QStringList defines = cmdLineParser.values(defineOption);
+            for (const QString &def : defines) {
+                const QStringList defs = def.split(QLatin1Char('='), QString::SkipEmptyParts);
+                if (!defs.isEmpty()) {
+                    preamble.append("#define");
+                    for (const QString &s : defs) {
+                        preamble.append(' ');
+                        preamble.append(s.toUtf8());
+                    }
+                    preamble.append('\n');
+                }
+            }
+            baker.setPreamble(preamble);
+        }
+
         QRhiShader bs = baker.bake();
         if (!bs.isValid()) {
             qWarning("Shader baking failed: %s", qPrintable(baker.errorMessage()));