Fix rewrite of inner queries in DisMaxQueryBuilder (#40956)

This commit implements missing rewrite for the DisMaxQueryBuilder.

Closes #40953

server/src/main/java/org/elasticsearch/index/query/DisMaxQueryBuilder.java

@@ -190,6 +190,27 @@ public class DisMaxQueryBuilder extends AbstractQueryBuilder<DisMaxQueryBuilder>
         return new DisjunctionMaxQuery(luceneQueries, tieBreaker);
     }
 
+    @Override
+    protected QueryBuilder doRewrite(QueryRewriteContext queryShardContext) throws IOException {
+        DisMaxQueryBuilder newBuilder = new DisMaxQueryBuilder();
+        boolean changed = false;
+        for (QueryBuilder query : queries) {
+            QueryBuilder result = query.rewrite(queryShardContext);
+            if (result != query) {
+                changed = true;
+            }
+            newBuilder.add(result);
+        }
+        if (changed) {
+            newBuilder.queryName(queryName);
+            newBuilder.boost(boost);
+            newBuilder.tieBreaker(tieBreaker);
+            return newBuilder;
+        } else {
+            return this;
+        }
+    }
+
     @Override
     protected int doHashCode() {
         return Objects.hash(queries, tieBreaker);

server/src/test/java/org/elasticsearch/index/query/DisMaxQueryBuilderTests.java

@@ -151,4 +151,19 @@ public class DisMaxQueryBuilderTests extends AbstractQueryTestCase<DisMaxQueryBu
         assertEquals(json, 0.7, parsed.tieBreaker(), 0.0001);
         assertEquals(json, 2, parsed.innerQueries().size());
     }
+
+    public void testRewriteMultipleTimes() throws IOException {
+        DisMaxQueryBuilder dismax = new DisMaxQueryBuilder();
+        dismax.add(new WrapperQueryBuilder(new WrapperQueryBuilder(new MatchAllQueryBuilder().toString()).toString()));
+        QueryBuilder rewritten = dismax.rewrite(createShardContext());
+        DisMaxQueryBuilder expected = new DisMaxQueryBuilder();
+        expected.add(new MatchAllQueryBuilder());
+        assertEquals(expected, rewritten);
+
+        expected = new DisMaxQueryBuilder();
+        expected.add(new MatchAllQueryBuilder());
+        QueryBuilder rewrittenAgain = rewritten.rewrite(createShardContext());
+        assertEquals(rewrittenAgain, expected);
+        assertEquals(Rewriteable.rewrite(dismax, createShardContext()), expected);
+    }
 }