From 71a5a2ec3ddac0b2d9c59d25a6b33f4fba182052 Mon Sep 17 00:00:00 2001
From: Florian Bernd <florian.bernd@elastic.co>
Date: Tue, 18 Jun 2024 13:03:06 +0200
Subject: [PATCH] Pass down type arguments to behaviors

---
 output/typescript/types.ts        |  6 +++---
 typescript-generator/src/index.ts | 36 ++++++++++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/output/typescript/types.ts b/output/typescript/types.ts
index 7afbfb827f..a3d8046957 100644
--- a/output/typescript/types.ts
+++ b/output/typescript/types.ts
@@ -5488,7 +5488,7 @@ export interface QueryDslConstantScoreQuery extends QueryDslQueryBase {
 export interface QueryDslDateDecayFunctionKeys extends QueryDslDecayFunctionBase<DateMath, Duration> {
 }
 export type QueryDslDateDecayFunction = QueryDslDateDecayFunctionKeys
-  & { [property: string]: QueryDslDecayPlacement<QueryDslTOrigin, QueryDslTScale> | QueryDslMultiValueMode }
+  & { [property: string]: QueryDslDecayPlacement<DateMath, Duration> | QueryDslMultiValueMode }
 
 export interface QueryDslDateDistanceFeatureQuery extends QueryDslDistanceFeatureQueryBase<DateMath, Duration> {
 }
@@ -5592,7 +5592,7 @@ export type QueryDslGeoBoundingBoxQuery = QueryDslGeoBoundingBoxQueryKeys
 export interface QueryDslGeoDecayFunctionKeys extends QueryDslDecayFunctionBase<GeoLocation, Distance> {
 }
 export type QueryDslGeoDecayFunction = QueryDslGeoDecayFunctionKeys
-  & { [property: string]: QueryDslDecayPlacement<QueryDslTOrigin, QueryDslTScale> | QueryDslMultiValueMode }
+  & { [property: string]: QueryDslDecayPlacement<GeoLocation, Distance> | QueryDslMultiValueMode }
 
 export interface QueryDslGeoDistanceFeatureQuery extends QueryDslDistanceFeatureQueryBase<GeoLocation, Distance> {
 }
@@ -5846,7 +5846,7 @@ export interface QueryDslNumberRangeQuery extends QueryDslRangeQueryBase<double>
 export interface QueryDslNumericDecayFunctionKeys extends QueryDslDecayFunctionBase<double, double> {
 }
 export type QueryDslNumericDecayFunction = QueryDslNumericDecayFunctionKeys
-  & { [property: string]: QueryDslDecayPlacement<QueryDslTOrigin, QueryDslTScale> | QueryDslMultiValueMode }
+  & { [property: string]: QueryDslDecayPlacement<double, double> | QueryDslMultiValueMode }
 
 export type QueryDslOperator = 'and' | 'AND' | 'or' | 'OR'
 
diff --git a/typescript-generator/src/index.ts b/typescript-generator/src/index.ts
index 1e6aa55eb8..6623f728e0 100644
--- a/typescript-generator/src/index.ts
+++ b/typescript-generator/src/index.ts
@@ -283,8 +283,21 @@ function buildBehaviorInterface (type: M.Interface): string {
 // The main difference with this approaches comes when you are actually using the types. 1 and 2 are good when
 // you are reading an object of that type, while 3 is good when you are writing an object of that type.
 function serializeAdditionalPropertiesType (type: M.Interface): string {
-  const dictionaryOf = lookupBehavior(type, 'AdditionalProperties') ?? lookupBehavior(type, 'AdditionalProperty')
+  let dictionaryOf = lookupBehavior(type, 'AdditionalProperties') ?? lookupBehavior(type, 'AdditionalProperty')
   if (dictionaryOf == null) throw new Error(`Unknown implements ${type.name.name}`)
+
+  const parent = model.types.find(t => t.name.name === type.inherits?.type.name)
+  if (parent != null && parent.kind === 'interface' && parent.generics != null && parent.generics.length === type.inherits?.generics?.length) {
+    const map = new Map<M.TypeName, M.ValueOf>();
+    for (let i = 0; i < parent.generics.length; i++) {
+      map.set(parent.generics[i], type.inherits.generics[i])
+    }
+    dictionaryOf = {
+      type: dictionaryOf.type,
+      generics: dictionaryOf.generics?.map(x => replaceGenerics(x, map))
+    }
+  }
+
   const extendedPropertyTypes = getAllExtendedPropertyTypes(type.inherits)
   const openGenerics = type.generics?.map(t => t.name) ?? []
   let code = `export interface ${createName(type.name)}Keys${buildGenerics(type.generics)}${buildInherits(type)} {\n`
@@ -325,6 +338,27 @@ function serializeAdditionalPropertiesType (type: M.Interface): string {
   }
 }
 
+function replaceGenerics(value: M.ValueOf, map: Map<M.TypeName, M.ValueOf>): M.ValueOf {
+  if (value.kind !== 'instance_of') {
+    return value
+  }
+
+  if (value.generics == null || value.generics.length === 0) {
+    for (const entry of map) {
+      if (entry[0].namespace === value.type.namespace && entry[0].name === value.type.name) {
+        return entry[1]
+      }
+    }
+    return value
+  }
+
+  return {
+    kind: 'instance_of',
+    type: value.type,
+    generics: value.generics.map(x => replaceGenerics(x, map))
+  }
+}
+
 function lookupBehavior (type: M.Interface, name: string): M.Inherits | null {
   if (!type.attachedBehaviors?.includes(name)) return null // eslint-disable-line
   if (Array.isArray(type.behaviors)) {