Adjust variable context to pass lezer-feel custom context tests

lib/zeebe/util/VariableContext.js

@@ -1,21 +1,29 @@
 import { VariableContext } from 'lezer-feel';
 
+
+/**
+ * @typedef { { entries?: Record<string, any>, [key: string]: any } } EntriesContextValue
+ */
+
+/**
+ * An alternative context that holds additional meta-data
+ */
 export class EntriesContext extends VariableContext {
+
+  /**
+   * @param {EntriesContextValue} value
+   */
   constructor(value = { entries: {} }) {
     super(value);
 
-    this.value.entries = this.value.entries || {};
-
-    const context = this.value;
-
-    for (const key in context.entries) {
-      const entry = context.entries[key];
+    const entries = this.value.entries = this.value.entries || {};
 
+    for (const [ key, entry ] of Object.entries(entries)) {
       if (entry instanceof EntriesContext) {
         continue;
       }
 
-      context.entries[key] = this.constructor.of(context.entries[key]);
+      entries[key] = EntriesContext.of(entry);
     }
   }
 
@@ -30,61 +38,78 @@ export class EntriesContext extends VariableContext {
       return value;
     }
 
-    if (value.atomic) {
-      return value.atomicValue;
+    const atomicValue = value?.value.atomicValue;
+
+    // keep value producer
+    if (atomicValue?.fn) {
+      return atomicValue;
     }
 
     return value;
   }
 
+  /**
+   * @param {string} key
+   * @param {any} value
+   *
+   * @return {this}
+   */
   set(key, value) {
-    return this.constructor.of(
+
+    const constructor = /** @type { typeof EntriesContext } */ (this.constructor);
+
+    return /** @type {this} */ (constructor.of(
       {
         ...this.value,
         entries: {
           ...this.value.entries,
           [key]: value
         }
       }
-    );
+    ));
   }
 
-  static of(...contexts) {
-    const unwrap = (context) => {
-
-      if (
-        this.isAtomic(context)
-      ) {
-        if (context instanceof this) {
-          return context.value;
-        }
+  /**
+   * @param { EntriesContext | EntriesContextValue | Record<string, any> } context
+   *
+   * @return { EntriesContextValue }
+   */
+  static __unwrap(context) {
+
+    if (this.isAtomic(context)) {
+      return context instanceof this
+        ? context.value
+        : { atomicValue: context };
+    }
 
-        return {
-          atomic: true,
-          atomicValue: context
-        };
-      }
+    return context;
+  }
 
-      return { ...context };
+  /**
+   * @param { EntriesContextValue } context
+   * @param { EntriesContextValue} other
+   *
+   * @return { EntriesContextValue }
+   */
+  static __merge(context, other) {
+
+    const {
+      entries: contextEntries = {},
+      ...contextRest
+    } = this.__unwrap(context);
+
+    const {
+      entries: otherEntries = {},
+      ...otherRest
+    } = this.__unwrap(other);
+
+    // @ts-ignore "access to internals"
+    const mergedEntries = super.__merge(contextEntries, otherEntries);
+
+    return {
+      ...contextRest,
+      ...otherRest,
+      entries: mergedEntries
     };
-
-    const merged = contexts.reduce((merged, context) => {
-
-      const {
-        entries = {},
-        ...rest
-      } = unwrap(context);
-
-      return {
-        ...merged,
-        ...rest,
-        entries: {
-          ...merged.entries,
-          ...entries
-        }
-      };
-    }, {});
-
-    return new this(merged);
   }
 }