From aac179031a193c3cc82bf29270e06ac6a26d2d72 Mon Sep 17 00:00:00 2001
From: pjanik <janikpiotrek@gmail.com>
Date: Tue, 23 Jul 2024 19:39:50 +0200
Subject: [PATCH 1/5] fix: update MathJS to v12.4.3 and handle the new approach
 to scope/partitioned map [PT-187950451]

---
 v3/package-lock.json                           | 18 +++++++++---------
 v3/package.json                                |  2 +-
 v3/src/models/formula/formula-types.ts         |  2 +-
 .../functions/aggregate-functions.test.ts      | 13 +++++++------
 .../formula/functions/aggregate-functions.ts   | 13 ++++++++-----
 .../formula/functions/lookup-functions.ts      |  6 ++++--
 v3/src/models/formula/functions/math.test.ts   | 12 ++++++++----
 v3/src/models/formula/functions/math.ts        | 15 +++++++++------
 .../functions/semi-aggregate-functions.ts      |  6 ++++--
 9 files changed, 51 insertions(+), 36 deletions(-)

diff --git a/v3/package-lock.json b/v3/package-lock.json
index 420461ac3f..a6f2724230 100644
--- a/v3/package-lock.json
+++ b/v3/package-lock.json
@@ -25,7 +25,7 @@
         "iframe-phone": "^1.3.1",
         "leaflet": "^1.9.4",
         "lodash": "^4.17.21",
-        "mathjs": "12.3.1",
+        "mathjs": "12.4.3",
         "mime": "^4.0.4",
         "mobx": "^6.13.0",
         "mobx-react-lite": "^4.0.7",
@@ -18339,11 +18339,11 @@
       }
     },
     "node_modules/mathjs": {
-      "version": "12.3.1",
-      "resolved": "https://registry.npmjs.org/mathjs/-/mathjs-12.3.1.tgz",
-      "integrity": "sha512-Z8G8Cunz97rV32eghW/twFFSY/dmU2eswnh6cAkcQMbbEUUpaweF0li435GkZjRS+G52VYotN+vLx3VTYBf5WA==",
+      "version": "12.4.3",
+      "resolved": "https://registry.npmjs.org/mathjs/-/mathjs-12.4.3.tgz",
+      "integrity": "sha512-oHdGPDbp7gO873xxG90RLq36IuicuKvbpr/bBG5g9c8Obm/VsKVrK9uoRZZHUodohzlnmCEqfDzbR3LH6m+aAQ==",
       "dependencies": {
-        "@babel/runtime": "^7.23.9",
+        "@babel/runtime": "^7.24.4",
         "complex.js": "^2.1.1",
         "decimal.js": "^10.4.3",
         "escape-latex": "^1.2.0",
@@ -38190,11 +38190,11 @@
       }
     },
     "mathjs": {
-      "version": "12.3.1",
-      "resolved": "https://registry.npmjs.org/mathjs/-/mathjs-12.3.1.tgz",
-      "integrity": "sha512-Z8G8Cunz97rV32eghW/twFFSY/dmU2eswnh6cAkcQMbbEUUpaweF0li435GkZjRS+G52VYotN+vLx3VTYBf5WA==",
+      "version": "12.4.3",
+      "resolved": "https://registry.npmjs.org/mathjs/-/mathjs-12.4.3.tgz",
+      "integrity": "sha512-oHdGPDbp7gO873xxG90RLq36IuicuKvbpr/bBG5g9c8Obm/VsKVrK9uoRZZHUodohzlnmCEqfDzbR3LH6m+aAQ==",
       "requires": {
-        "@babel/runtime": "^7.23.9",
+        "@babel/runtime": "^7.24.4",
         "complex.js": "^2.1.1",
         "decimal.js": "^10.4.3",
         "escape-latex": "^1.2.0",
diff --git a/v3/package.json b/v3/package.json
index f75179560d..57154beb46 100644
--- a/v3/package.json
+++ b/v3/package.json
@@ -191,7 +191,7 @@
     "iframe-phone": "^1.3.1",
     "leaflet": "^1.9.4",
     "lodash": "^4.17.21",
-    "mathjs": "12.3.1",
+    "mathjs": "12.4.3",
     "mime": "^4.0.4",
     "mobx": "^6.13.0",
     "mobx-react-lite": "^4.0.7",
diff --git a/v3/src/models/formula/formula-types.ts b/v3/src/models/formula/formula-types.ts
index a9487686aa..3e8eb3d7f0 100644
--- a/v3/src/models/formula/formula-types.ts
+++ b/v3/src/models/formula/formula-types.ts
@@ -41,7 +41,7 @@ export type EvaluateFunc = (...args: FValue[]) => FValue
 
 export type EvaluateFuncWithAggregateContextSupport = (...args: (FValueOrArray)[]) => FValueOrArray
 
-export type EvaluateRawFunc = (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => FValueOrArray
+export type EvaluateRawFunc = (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => FValueOrArray
 
 export type CaseList = ICase[] | "ALL_CASES"
 
diff --git a/v3/src/models/formula/functions/aggregate-functions.test.ts b/v3/src/models/formula/functions/aggregate-functions.test.ts
index 52f64cc5fa..e1b4338ce2 100644
--- a/v3/src/models/formula/functions/aggregate-functions.test.ts
+++ b/v3/src/models/formula/functions/aggregate-functions.test.ts
@@ -151,25 +151,26 @@ describe("cachedAggregateFnFactory", () => {
       setCached: (key: string, val: any) => cache.set(key, val),
       getCached: (key: string) => cache.get(key)
     } as any as FormulaMathJsScope
+    const partitionedMap = { a: scope }
 
-    expect(cachedFn([ parse("1"), parse("2") ], null, scope)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2") ], null, partitionedMap)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(1)
 
-    expect(cachedFn([ parse("1"), parse("2") ], null, scope)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2") ], null, partitionedMap)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(1) // Same arguments as in the previous call, cache should be used.
 
-    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, scope)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, partitionedMap)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(2) // Different arguments, so cache should not be used.
 
-    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, scope)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, partitionedMap)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(2) // Same arguments as in the previous call, cache should be used.
 
     // Update scope.getCaseAggregateGroupId
     ;(scope.getCaseAggregateGroupId as jest.Mock).mockImplementation(() => "newTestGroup")
-    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, scope)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, partitionedMap)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(3) // New caseAggregateGroupId, so cache should not be used.
 
-    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, scope)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, partitionedMap)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(3) // Same arguments and caseAggregateGroupId, cache should be used.
   })
 })
diff --git a/v3/src/models/formula/functions/aggregate-functions.ts b/v3/src/models/formula/functions/aggregate-functions.ts
index 9b67752490..c1ea9dbced 100644
--- a/v3/src/models/formula/functions/aggregate-functions.ts
+++ b/v3/src/models/formula/functions/aggregate-functions.ts
@@ -5,14 +5,15 @@ import { UNDEF_RESULT, evaluateNode, isNumber, isValueTruthy } from "./function-
 
 // Almost every aggregate function can be cached in the same way.
 export const cachedAggregateFnFactory =
-(fnName: string, fn: (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => FValueOrArray) => {
-  return (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+(fnName: string, fn: (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => FValueOrArray) => {
+  return (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+    const scope = partitionedMap.a
     const cacheKey = `${fnName}(${args.toString()})-${scope.getCaseAggregateGroupId()}`
     const cachedValue = scope.getCached(cacheKey)
     if (cachedValue !== undefined) {
       return cachedValue
     }
-    const result = fn(args, mathjs, scope)
+    const result = fn(args, mathjs, partitionedMap)
     scope.setCached(cacheKey, result)
     return result
   }
@@ -21,7 +22,8 @@ export const cachedAggregateFnFactory =
 // Note that aggregate functions like mean, max, min, etc., all have exactly the same signature and implementation.
 // The only difference is the final math operation applies to the expression results.
 const aggregateFnWithFilterFactory = (fn: (values: number[]) => FValue) => {
-  return (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+  return (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+    const scope = partitionedMap.a
     const [ expression, filter ] = args
     let expressionValues = evaluateNode(expression, scope)
     if (!Array.isArray(expressionValues)) {
@@ -96,7 +98,8 @@ export const aggregateFunctions = {
     // arguments, the default caching method would calculate incorrect cache key. Hence, caching is implemented directly
     // in the function body.
     cachedEvaluateFactory: undefined,
-    evaluateRaw: (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+      const scope = partitionedMap.a
       const [expression, filter] = args
       if (!expression) {
         // Special case: count() without arguments returns number of children cases. Note that this cannot be cached
diff --git a/v3/src/models/formula/functions/lookup-functions.ts b/v3/src/models/formula/functions/lookup-functions.ts
index 1732f47f70..649fc2cfce 100644
--- a/v3/src/models/formula/functions/lookup-functions.ts
+++ b/v3/src/models/formula/functions/lookup-functions.ts
@@ -32,7 +32,8 @@ export const lookupFunctions = {
         attrNameArg.value = displayNameMap.dataSet[dataSetName]?.attribute[attrName] || attrName
       }
     },
-    evaluateRaw: (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+      const scope = partitionedMap.a
       const functionName = "lookupByIndex"
       const numOfReqArgs = lookupFunctions.lookupByIndex.numOfRequiredArguments
       if (args.length !== numOfReqArgs) {
@@ -91,7 +92,8 @@ export const lookupFunctions = {
         keyAttrNameArg.value = displayNameMap.dataSet[dataSetName]?.attribute[keyAttrName] || keyAttrName
       }
     },
-    evaluateRaw: (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+      const scope = partitionedMap.a
       const functionName = "lookupByKey"
       const numOfReqArgs = lookupFunctions.lookupByKey.numOfRequiredArguments
       if (args.length !== numOfReqArgs) {
diff --git a/v3/src/models/formula/functions/math.test.ts b/v3/src/models/formula/functions/math.test.ts
index 8515256564..5dbd9887e9 100644
--- a/v3/src/models/formula/functions/math.test.ts
+++ b/v3/src/models/formula/functions/math.test.ts
@@ -17,10 +17,11 @@ describe("evaluateRawWithAggregateContext", () => {
         scope.aggregateContext = false
       }
     }
+    const partitionedMap = { a: scope } as any as { a: FormulaMathJsScope }
     let result = false
     const mockFn = jest.fn(() => { result = scope.aggregateContext; return "" })
 
-    evaluateRawWithAggregateContext(mockFn)(args, mathjs, scope as any as FormulaMathJsScope)
+    evaluateRawWithAggregateContext(mockFn)(args, mathjs, partitionedMap)
     expect(mockFn).toHaveBeenCalledWith(args, mathjs, scope)
     expect(result).toBeTruthy()
   })
@@ -33,10 +34,11 @@ describe("evaluateRawWithDefaultArg", () => {
     const scope = {
       defaultArgumentNode: { name: "default" }
     }
+    const partitionedMap = { a: scope } as any as { a: FormulaMathJsScope }
     const mockFn = jest.fn((_args: MathNode[]) => (_args[0] as SymbolNode).name)
 
     const numOfReqArgs = 1
-    const res = evaluateRawWithDefaultArg(mockFn, numOfReqArgs)(args, mathjs, scope as any as FormulaMathJsScope)
+    const res = evaluateRawWithDefaultArg(mockFn, numOfReqArgs)(args, mathjs, partitionedMap)
     expect(mockFn).toHaveBeenCalledWith([scope.defaultArgumentNode], mathjs, scope)
     expect(res).toEqual("default")
   })
@@ -47,10 +49,11 @@ describe("evaluateRawWithDefaultArg", () => {
     const scope = {
       defaultArgumentNode: { name: "default" }
     }
+    const partitionedMap = { a: scope } as any as { a: FormulaMathJsScope }
     const mockFn = jest.fn((_args: MathNode[]) => (_args[0] as SymbolNode).name)
 
     const res =
-      evaluateRawWithDefaultArg(mockFn, 1)(args as any as MathNode[], mathjs, scope as any as FormulaMathJsScope)
+      evaluateRawWithDefaultArg(mockFn, 1)(args as any as MathNode[], mathjs, partitionedMap)
     expect(mockFn).toHaveBeenCalledWith(args, mathjs, scope)
     expect(res).toEqual("provided")
   })
@@ -61,9 +64,10 @@ describe("evaluateToEvaluateRaw", () => {
     const args = [ parse("1"), parse("2") ]
     const mathjs = {}
     const scope = {}
+    const partitionedMap = { a: scope } as any as { a: FormulaMathJsScope }
     const mockFn = jest.fn((a: FValueOrArray, b: FValueOrArray) => Number(a) + Number(b))
 
-    const res = evaluateToEvaluateRaw(mockFn)(args as any as MathNode[], mathjs, scope as any as FormulaMathJsScope)
+    const res = evaluateToEvaluateRaw(mockFn)(args as any as MathNode[], mathjs, partitionedMap)
     expect(mockFn).toHaveBeenCalledWith(1, 2)
     expect(res).toEqual(3)
   })
diff --git a/v3/src/models/formula/functions/math.ts b/v3/src/models/formula/functions/math.ts
index b52f2429fe..262e6f9685 100644
--- a/v3/src/models/formula/functions/math.ts
+++ b/v3/src/models/formula/functions/math.ts
@@ -18,23 +18,26 @@ export const math = create(all)
 
 // Each aggregate function needs to be evaluated with `withAggregateContext` method.
 export const evaluateRawWithAggregateContext = (fn: EvaluateRawFunc): EvaluateRawFunc => {
-  return (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+  return (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+    const scope = partitionedMap.a
     // withAggregateContext returns result of the callback function
-    return scope.withAggregateContext(() => fn(args, mathjs, scope))
+    return scope.withAggregateContext(() => fn(args, mathjs, partitionedMap))
   }
 }
 
 export const evaluateRawWithDefaultArg = (fn: EvaluateRawFunc, numOfRequiredArguments: number): EvaluateRawFunc => {
-  return (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+  return (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+    const scope = partitionedMap.a
     if (scope.defaultArgumentNode && args.length < numOfRequiredArguments) {
-      return fn([...args, scope.defaultArgumentNode], mathjs, scope)
+      return fn([...args, scope.defaultArgumentNode], mathjs, partitionedMap)
     }
-    return fn(args, mathjs, scope)
+    return fn(args, mathjs, partitionedMap)
   }
 }
 
 export const evaluateToEvaluateRaw = (fn: EvaluateFuncWithAggregateContextSupport): EvaluateRawFunc => {
-  return (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+  return (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+    const scope = partitionedMap.a
     return fn(...(args.map(arg => evaluateNode(arg, scope))))
   }
 }
diff --git a/v3/src/models/formula/functions/semi-aggregate-functions.ts b/v3/src/models/formula/functions/semi-aggregate-functions.ts
index 774d013648..e940183c5f 100644
--- a/v3/src/models/formula/functions/semi-aggregate-functions.ts
+++ b/v3/src/models/formula/functions/semi-aggregate-functions.ts
@@ -8,12 +8,13 @@ export const semiAggregateFunctions = {
     numOfRequiredArguments: 1,
     // expression and filter are evaluated as aggregate symbols, defaultValue is not - it depends on case index
     isSemiAggregate: [true, false, true],
-    evaluateRaw: (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
       interface ICachedData {
         result?: FValue
         resultCasePointer: number
       }
 
+      const scope = partitionedMap.a
       const caseGroupId = scope.getCaseGroupId()
       const cacheKey = `next(${args.toString()})-${caseGroupId}`
       const [ expression, defaultValue, filter ] = args
@@ -75,7 +76,8 @@ export const semiAggregateFunctions = {
     selfReferenceAllowed: true,
     // expression and filter are evaluated as aggregate symbols, defaultValue is not - it depends on case index
     isSemiAggregate: [true, false, true],
-    evaluateRaw: (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+    evaluateRaw: (args: MathNode[], mathjs: any, parentId: { a: FormulaMathJsScope }) => {
+      const scope = parentId.a
       const [ expression, defaultValue, filter ] = args
 
       const caseGroupId = scope.getCaseGroupId()

From 31d9d705d16bde98ce34f03aa6936bea60ef5542 Mon Sep 17 00:00:00 2001
From: pjanik <janikpiotrek@gmail.com>
Date: Wed, 24 Jul 2024 12:33:30 +0200
Subject: [PATCH 2/5] chore: add MathJSPartitionedMap type [PT-187950451]

---
 v3/src/models/formula/formula-mathjs-scope.ts   |  8 --------
 v3/src/models/formula/formula-types.ts          |  7 ++++++-
 .../formula/functions/aggregate-functions.ts    | 11 +++++------
 .../formula/functions/lookup-functions.ts       |  7 +++----
 v3/src/models/formula/functions/math.test.ts    | 17 ++++++++---------
 v3/src/models/formula/functions/math.ts         |  9 ++++-----
 .../functions/semi-aggregate-functions.ts       |  9 ++++-----
 7 files changed, 30 insertions(+), 38 deletions(-)

diff --git a/v3/src/models/formula/formula-mathjs-scope.ts b/v3/src/models/formula/formula-mathjs-scope.ts
index 44f2bffad8..c8ad54cb5f 100644
--- a/v3/src/models/formula/formula-mathjs-scope.ts
+++ b/v3/src/models/formula/formula-mathjs-scope.ts
@@ -129,14 +129,6 @@ export class FormulaMathJsScope {
     throw new Error("It's not allowed to clear values in the formula scope.")
   }
 
-  // MathJS requests a separate subscope for every parsed and evaluated function call.
-  createSubScope () {
-    // In regular MathJS use case, a subscope is used to create a new scope for a function call. It needs to ensure
-    // that variables from subscope cannot overwrite variables from the parent scope. However, since we don't allow
-    // any variables to be set in the formula scope, we can reuse the same scope instance.
-    return this
-  }
-
   // --- Custom functions used by our formulas or formula manager --
 
   setExtraScope(extraScope?: Map<string, FValue>) {
diff --git a/v3/src/models/formula/formula-types.ts b/v3/src/models/formula/formula-types.ts
index 3e8eb3d7f0..f483b8b317 100644
--- a/v3/src/models/formula/formula-types.ts
+++ b/v3/src/models/formula/formula-types.ts
@@ -41,7 +41,12 @@ export type EvaluateFunc = (...args: FValue[]) => FValue
 
 export type EvaluateFuncWithAggregateContextSupport = (...args: (FValueOrArray)[]) => FValueOrArray
 
-export type EvaluateRawFunc = (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => FValueOrArray
+// MathJS v12.3.2 introduced the concept of PartitionedMap, replacing the previous concepts of scopes and sub-scopes.
+// The parent scope is still available as the `a` property of the PartitionedMap.
+// For more information, see: https://github.com/josdejong/mathjs/pull/3150
+export type MathJSPartitionedMap = { a: FormulaMathJsScope }
+
+export type EvaluateRawFunc = (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => FValueOrArray
 
 export type CaseList = ICase[] | "ALL_CASES"
 
diff --git a/v3/src/models/formula/functions/aggregate-functions.ts b/v3/src/models/formula/functions/aggregate-functions.ts
index c1ea9dbced..c95ca659fc 100644
--- a/v3/src/models/formula/functions/aggregate-functions.ts
+++ b/v3/src/models/formula/functions/aggregate-functions.ts
@@ -1,12 +1,11 @@
 import { MathNode, mad, max, mean, median, min, sum } from "mathjs"
-import { FormulaMathJsScope } from "../formula-mathjs-scope"
-import { FValue, FValueOrArray } from "../formula-types"
+import { FValue, FValueOrArray, MathJSPartitionedMap } from "../formula-types"
 import { UNDEF_RESULT, evaluateNode, isNumber, isValueTruthy } from "./function-utils"
 
 // Almost every aggregate function can be cached in the same way.
 export const cachedAggregateFnFactory =
-(fnName: string, fn: (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => FValueOrArray) => {
-  return (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+(fnName: string, fn: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => FValueOrArray) => {
+  return (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
     const scope = partitionedMap.a
     const cacheKey = `${fnName}(${args.toString()})-${scope.getCaseAggregateGroupId()}`
     const cachedValue = scope.getCached(cacheKey)
@@ -22,7 +21,7 @@ export const cachedAggregateFnFactory =
 // Note that aggregate functions like mean, max, min, etc., all have exactly the same signature and implementation.
 // The only difference is the final math operation applies to the expression results.
 const aggregateFnWithFilterFactory = (fn: (values: number[]) => FValue) => {
-  return (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+  return (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
     const scope = partitionedMap.a
     const [ expression, filter ] = args
     let expressionValues = evaluateNode(expression, scope)
@@ -98,7 +97,7 @@ export const aggregateFunctions = {
     // arguments, the default caching method would calculate incorrect cache key. Hence, caching is implemented directly
     // in the function body.
     cachedEvaluateFactory: undefined,
-    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
       const scope = partitionedMap.a
       const [expression, filter] = args
       if (!expression) {
diff --git a/v3/src/models/formula/functions/lookup-functions.ts b/v3/src/models/formula/functions/lookup-functions.ts
index 649fc2cfce..7596689402 100644
--- a/v3/src/models/formula/functions/lookup-functions.ts
+++ b/v3/src/models/formula/functions/lookup-functions.ts
@@ -1,6 +1,5 @@
 import { ConstantNode, MathNode } from "mathjs"
-import { FormulaMathJsScope } from "../formula-mathjs-scope"
-import { DisplayNameMap, FValue, ILookupDependency } from "../formula-types"
+import { DisplayNameMap, FValue, ILookupDependency, MathJSPartitionedMap } from "../formula-types"
 import { rmCanonicalPrefix } from "../utils/name-mapping-utils"
 import { UNDEF_RESULT, equal, evaluateNode } from "./function-utils"
 import { isConstantStringNode } from "../utils/mathjs-utils"
@@ -32,7 +31,7 @@ export const lookupFunctions = {
         attrNameArg.value = displayNameMap.dataSet[dataSetName]?.attribute[attrName] || attrName
       }
     },
-    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
       const scope = partitionedMap.a
       const functionName = "lookupByIndex"
       const numOfReqArgs = lookupFunctions.lookupByIndex.numOfRequiredArguments
@@ -92,7 +91,7 @@ export const lookupFunctions = {
         keyAttrNameArg.value = displayNameMap.dataSet[dataSetName]?.attribute[keyAttrName] || keyAttrName
       }
     },
-    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
       const scope = partitionedMap.a
       const functionName = "lookupByKey"
       const numOfReqArgs = lookupFunctions.lookupByKey.numOfRequiredArguments
diff --git a/v3/src/models/formula/functions/math.test.ts b/v3/src/models/formula/functions/math.test.ts
index 5dbd9887e9..1aa69da8f8 100644
--- a/v3/src/models/formula/functions/math.test.ts
+++ b/v3/src/models/formula/functions/math.test.ts
@@ -1,9 +1,8 @@
 import { MathNode, SymbolNode, parse } from "mathjs"
-import { FormulaMathJsScope } from "../formula-mathjs-scope"
 import {
   evaluateRawWithAggregateContext, evaluateRawWithDefaultArg, evaluateToEvaluateRaw, evaluateWithAggregateContextSupport
 } from "./math"
-import { FValue, FValueOrArray } from "../formula-types"
+import { FValue, FValueOrArray, MathJSPartitionedMap } from "../formula-types"
 
 describe("evaluateRawWithAggregateContext", () => {
   it("should call provided function within withAggregateContext", () => {
@@ -17,12 +16,12 @@ describe("evaluateRawWithAggregateContext", () => {
         scope.aggregateContext = false
       }
     }
-    const partitionedMap = { a: scope } as any as { a: FormulaMathJsScope }
+    const partitionedMap = { a: scope } as any as MathJSPartitionedMap
     let result = false
     const mockFn = jest.fn(() => { result = scope.aggregateContext; return "" })
 
     evaluateRawWithAggregateContext(mockFn)(args, mathjs, partitionedMap)
-    expect(mockFn).toHaveBeenCalledWith(args, mathjs, scope)
+    expect(mockFn).toHaveBeenCalledWith(args, mathjs, partitionedMap)
     expect(result).toBeTruthy()
   })
 })
@@ -34,12 +33,12 @@ describe("evaluateRawWithDefaultArg", () => {
     const scope = {
       defaultArgumentNode: { name: "default" }
     }
-    const partitionedMap = { a: scope } as any as { a: FormulaMathJsScope }
+    const partitionedMap = { a: scope } as any as MathJSPartitionedMap
     const mockFn = jest.fn((_args: MathNode[]) => (_args[0] as SymbolNode).name)
 
     const numOfReqArgs = 1
     const res = evaluateRawWithDefaultArg(mockFn, numOfReqArgs)(args, mathjs, partitionedMap)
-    expect(mockFn).toHaveBeenCalledWith([scope.defaultArgumentNode], mathjs, scope)
+    expect(mockFn).toHaveBeenCalledWith([scope.defaultArgumentNode], mathjs, partitionedMap)
     expect(res).toEqual("default")
   })
 
@@ -49,12 +48,12 @@ describe("evaluateRawWithDefaultArg", () => {
     const scope = {
       defaultArgumentNode: { name: "default" }
     }
-    const partitionedMap = { a: scope } as any as { a: FormulaMathJsScope }
+    const partitionedMap = { a: scope } as any as MathJSPartitionedMap
     const mockFn = jest.fn((_args: MathNode[]) => (_args[0] as SymbolNode).name)
 
     const res =
       evaluateRawWithDefaultArg(mockFn, 1)(args as any as MathNode[], mathjs, partitionedMap)
-    expect(mockFn).toHaveBeenCalledWith(args, mathjs, scope)
+    expect(mockFn).toHaveBeenCalledWith(args, mathjs, partitionedMap)
     expect(res).toEqual("provided")
   })
 })
@@ -64,7 +63,7 @@ describe("evaluateToEvaluateRaw", () => {
     const args = [ parse("1"), parse("2") ]
     const mathjs = {}
     const scope = {}
-    const partitionedMap = { a: scope } as any as { a: FormulaMathJsScope }
+    const partitionedMap = { a: scope } as any as MathJSPartitionedMap
     const mockFn = jest.fn((a: FValueOrArray, b: FValueOrArray) => Number(a) + Number(b))
 
     const res = evaluateToEvaluateRaw(mockFn)(args as any as MathNode[], mathjs, partitionedMap)
diff --git a/v3/src/models/formula/functions/math.ts b/v3/src/models/formula/functions/math.ts
index 262e6f9685..55b76965a4 100644
--- a/v3/src/models/formula/functions/math.ts
+++ b/v3/src/models/formula/functions/math.ts
@@ -1,8 +1,7 @@
 import { create, all, MathNode } from 'mathjs'
-import { FormulaMathJsScope } from '../formula-mathjs-scope'
 import {
   CODAPMathjsFunctionRegistry, EvaluateFunc, EvaluateFuncWithAggregateContextSupport, EvaluateRawFunc, FValue,
-  FValueOrArray
+  FValueOrArray, MathJSPartitionedMap
 } from '../formula-types'
 import { evaluateNode } from './function-utils'
 import { arithmeticFunctions } from './arithmetic-functions'
@@ -18,7 +17,7 @@ export const math = create(all)
 
 // Each aggregate function needs to be evaluated with `withAggregateContext` method.
 export const evaluateRawWithAggregateContext = (fn: EvaluateRawFunc): EvaluateRawFunc => {
-  return (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+  return (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
     const scope = partitionedMap.a
     // withAggregateContext returns result of the callback function
     return scope.withAggregateContext(() => fn(args, mathjs, partitionedMap))
@@ -26,7 +25,7 @@ export const evaluateRawWithAggregateContext = (fn: EvaluateRawFunc): EvaluateRa
 }
 
 export const evaluateRawWithDefaultArg = (fn: EvaluateRawFunc, numOfRequiredArguments: number): EvaluateRawFunc => {
-  return (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+  return (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
     const scope = partitionedMap.a
     if (scope.defaultArgumentNode && args.length < numOfRequiredArguments) {
       return fn([...args, scope.defaultArgumentNode], mathjs, partitionedMap)
@@ -36,7 +35,7 @@ export const evaluateRawWithDefaultArg = (fn: EvaluateRawFunc, numOfRequiredArgu
 }
 
 export const evaluateToEvaluateRaw = (fn: EvaluateFuncWithAggregateContextSupport): EvaluateRawFunc => {
-  return (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+  return (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
     const scope = partitionedMap.a
     return fn(...(args.map(arg => evaluateNode(arg, scope))))
   }
diff --git a/v3/src/models/formula/functions/semi-aggregate-functions.ts b/v3/src/models/formula/functions/semi-aggregate-functions.ts
index e940183c5f..9a47c9e774 100644
--- a/v3/src/models/formula/functions/semi-aggregate-functions.ts
+++ b/v3/src/models/formula/functions/semi-aggregate-functions.ts
@@ -1,6 +1,5 @@
 import { MathNode } from "mathjs"
-import { FormulaMathJsScope } from "../formula-mathjs-scope"
-import { FValue } from "../formula-types"
+import { FValue, MathJSPartitionedMap } from "../formula-types"
 import { UNDEF_RESULT, evaluateNode, isValueTruthy } from "./function-utils"
 
 export const semiAggregateFunctions = {
@@ -8,7 +7,7 @@ export const semiAggregateFunctions = {
     numOfRequiredArguments: 1,
     // expression and filter are evaluated as aggregate symbols, defaultValue is not - it depends on case index
     isSemiAggregate: [true, false, true],
-    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: { a: FormulaMathJsScope }) => {
+    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
       interface ICachedData {
         result?: FValue
         resultCasePointer: number
@@ -76,8 +75,8 @@ export const semiAggregateFunctions = {
     selfReferenceAllowed: true,
     // expression and filter are evaluated as aggregate symbols, defaultValue is not - it depends on case index
     isSemiAggregate: [true, false, true],
-    evaluateRaw: (args: MathNode[], mathjs: any, parentId: { a: FormulaMathJsScope }) => {
-      const scope = parentId.a
+    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
+      const scope = partitionedMap.a
       const [ expression, defaultValue, filter ] = args
 
       const caseGroupId = scope.getCaseGroupId()

From 30cd1c48af79122d26d8ba7326a5532968c5213d Mon Sep 17 00:00:00 2001
From: pjanik <janikpiotrek@gmail.com>
Date: Wed, 24 Jul 2024 19:02:21 +0200
Subject: [PATCH 3/5] fix: refactor/improve handling of MathJS's PartitionedMap
 [PT-187950451]

---
 v3/src/models/formula/formula-types.ts        |  6 ++-
 .../functions/aggregate-functions.test.ts     | 14 +++----
 .../formula/functions/aggregate-functions.ts  | 20 +++++-----
 .../formula/functions/function-utils.test.ts  | 40 ++++++++++++++++++-
 .../formula/functions/function-utils.ts       | 21 ++++++++++
 .../formula/functions/lookup-functions.ts     | 12 +++---
 v3/src/models/formula/functions/math.test.ts  | 22 +++++-----
 v3/src/models/formula/functions/math.ts       | 24 +++++------
 .../formula/functions/other-functions.ts      |  2 +
 .../functions/semi-aggregate-functions.ts     | 12 +++---
 10 files changed, 118 insertions(+), 55 deletions(-)

diff --git a/v3/src/models/formula/formula-types.ts b/v3/src/models/formula/formula-types.ts
index f483b8b317..d19217d78b 100644
--- a/v3/src/models/formula/formula-types.ts
+++ b/v3/src/models/formula/formula-types.ts
@@ -44,9 +44,11 @@ export type EvaluateFuncWithAggregateContextSupport = (...args: (FValueOrArray)[
 // MathJS v12.3.2 introduced the concept of PartitionedMap, replacing the previous concepts of scopes and sub-scopes.
 // The parent scope is still available as the `a` property of the PartitionedMap.
 // For more information, see: https://github.com/josdejong/mathjs/pull/3150
-export type MathJSPartitionedMap = { a: FormulaMathJsScope }
+export type MathJSPartitionedMap = { a: CurrentScope, b: Map<string, any>}
 
-export type EvaluateRawFunc = (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => FValueOrArray
+export type CurrentScope = MathJSPartitionedMap | FormulaMathJsScope
+
+export type EvaluateRawFunc = (args: MathNode[], mathjs: any, currentScope: CurrentScope) => FValueOrArray
 
 export type CaseList = ICase[] | "ALL_CASES"
 
diff --git a/v3/src/models/formula/functions/aggregate-functions.test.ts b/v3/src/models/formula/functions/aggregate-functions.test.ts
index e1b4338ce2..eef56aa868 100644
--- a/v3/src/models/formula/functions/aggregate-functions.test.ts
+++ b/v3/src/models/formula/functions/aggregate-functions.test.ts
@@ -151,26 +151,26 @@ describe("cachedAggregateFnFactory", () => {
       setCached: (key: string, val: any) => cache.set(key, val),
       getCached: (key: string) => cache.get(key)
     } as any as FormulaMathJsScope
-    const partitionedMap = { a: scope }
+    const currentScope = { a: scope, b: new Map() }
 
-    expect(cachedFn([ parse("1"), parse("2") ], null, partitionedMap)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2") ], null, currentScope)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(1)
 
-    expect(cachedFn([ parse("1"), parse("2") ], null, partitionedMap)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2") ], null, currentScope)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(1) // Same arguments as in the previous call, cache should be used.
 
-    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, partitionedMap)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, currentScope)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(2) // Different arguments, so cache should not be used.
 
-    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, partitionedMap)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, currentScope)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(2) // Same arguments as in the previous call, cache should be used.
 
     // Update scope.getCaseAggregateGroupId
     ;(scope.getCaseAggregateGroupId as jest.Mock).mockImplementation(() => "newTestGroup")
-    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, partitionedMap)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, currentScope)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(3) // New caseAggregateGroupId, so cache should not be used.
 
-    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, partitionedMap)).toEqual(123)
+    expect(cachedFn([ parse("1"), parse("2"), parse("3") ], null, currentScope)).toEqual(123)
     expect(fn).toHaveBeenCalledTimes(3) // Same arguments and caseAggregateGroupId, cache should be used.
   })
 })
diff --git a/v3/src/models/formula/functions/aggregate-functions.ts b/v3/src/models/formula/functions/aggregate-functions.ts
index c95ca659fc..111f9ca6a4 100644
--- a/v3/src/models/formula/functions/aggregate-functions.ts
+++ b/v3/src/models/formula/functions/aggregate-functions.ts
@@ -1,18 +1,18 @@
 import { MathNode, mad, max, mean, median, min, sum } from "mathjs"
-import { FValue, FValueOrArray, MathJSPartitionedMap } from "../formula-types"
-import { UNDEF_RESULT, evaluateNode, isNumber, isValueTruthy } from "./function-utils"
+import { CurrentScope, FValue, FValueOrArray } from "../formula-types"
+import { UNDEF_RESULT, evaluateNode, getRootScope, isNumber, isValueTruthy } from "./function-utils"
 
 // Almost every aggregate function can be cached in the same way.
 export const cachedAggregateFnFactory =
-(fnName: string, fn: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => FValueOrArray) => {
-  return (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
-    const scope = partitionedMap.a
+(fnName: string, fn: (args: MathNode[], mathjs: any, currentScope: CurrentScope) => FValueOrArray) => {
+  return (args: MathNode[], mathjs: any, currentScope: CurrentScope) => {
+    const scope = getRootScope(currentScope)
     const cacheKey = `${fnName}(${args.toString()})-${scope.getCaseAggregateGroupId()}`
     const cachedValue = scope.getCached(cacheKey)
     if (cachedValue !== undefined) {
       return cachedValue
     }
-    const result = fn(args, mathjs, partitionedMap)
+    const result = fn(args, mathjs, currentScope)
     scope.setCached(cacheKey, result)
     return result
   }
@@ -21,8 +21,8 @@ export const cachedAggregateFnFactory =
 // Note that aggregate functions like mean, max, min, etc., all have exactly the same signature and implementation.
 // The only difference is the final math operation applies to the expression results.
 const aggregateFnWithFilterFactory = (fn: (values: number[]) => FValue) => {
-  return (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
-    const scope = partitionedMap.a
+  return (args: MathNode[], mathjs: any, currentScope: CurrentScope) => {
+    const scope = getRootScope(currentScope)
     const [ expression, filter ] = args
     let expressionValues = evaluateNode(expression, scope)
     if (!Array.isArray(expressionValues)) {
@@ -97,8 +97,8 @@ export const aggregateFunctions = {
     // arguments, the default caching method would calculate incorrect cache key. Hence, caching is implemented directly
     // in the function body.
     cachedEvaluateFactory: undefined,
-    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
-      const scope = partitionedMap.a
+    evaluateRaw: (args: MathNode[], mathjs: any, currentScope: CurrentScope) => {
+      const scope = getRootScope(currentScope)
       const [expression, filter] = args
       if (!expression) {
         // Special case: count() without arguments returns number of children cases. Note that this cannot be cached
diff --git a/v3/src/models/formula/functions/function-utils.test.ts b/v3/src/models/formula/functions/function-utils.test.ts
index b21779a826..c17d2113bf 100644
--- a/v3/src/models/formula/functions/function-utils.test.ts
+++ b/v3/src/models/formula/functions/function-utils.test.ts
@@ -1,5 +1,7 @@
-import { isValueTruthy, equal, evaluateNode } from "./function-utils"
+import { isValueTruthy, equal, evaluateNode, isPartitionedMap, isMap, getRootScope } from "./function-utils"
 import { parse } from "mathjs"
+import { CurrentScope } from "../formula-types"
+import { FormulaMathJsScope } from "../formula-mathjs-scope"
 
 describe("isValueTruthy", () => {
   it("should return true for truthy values", () => {
@@ -42,3 +44,39 @@ describe("evaluateNode", () => {
     expect(evaluateNode(node)).toBe(3)
   })
 })
+
+describe("isPartitionedMap", () => {
+  it("should return true for partitioned map", () => {
+    const map = { a: {}, b: new Map() }
+    expect(isPartitionedMap(map)).toBe(true)
+  })
+
+  it("should return false for non-partitioned map", () => {
+    const map = new Map()
+    expect(isPartitionedMap(map)).toBe(false)
+  })
+})
+
+describe("isMap", () => {
+  it("should return true for map", () => {
+    const map = new Map()
+    expect(isMap(map)).toBe(true)
+  })
+
+  it("should return false for non-map", () => {
+    const map = {}
+    expect(isMap(map)).toBe(false)
+  })
+})
+
+describe("getRootScope", () => {
+  it("should return root scope for partitioned map", () => {
+    const scope = { a: {}, b: new Map() }
+    expect(getRootScope(scope as any as CurrentScope)).toBe(scope.a)
+  })
+
+  it("should return scope for non-partitioned map", () => {
+    const scope = {} as any as FormulaMathJsScope
+    expect(getRootScope(scope)).toBe(scope)
+  })
+})
diff --git a/v3/src/models/formula/functions/function-utils.ts b/v3/src/models/formula/functions/function-utils.ts
index aa4a811fbe..9bba84d34e 100644
--- a/v3/src/models/formula/functions/function-utils.ts
+++ b/v3/src/models/formula/functions/function-utils.ts
@@ -1,6 +1,7 @@
 import { MathNode } from "mathjs"
 import { FormulaMathJsScope } from "../formula-mathjs-scope"
 import { isValueNonEmpty } from "../../../utilities/math-utils"
+import { CurrentScope, MathJSPartitionedMap } from "../formula-types"
 export { isNumber, isValueNonEmpty } from "../../../utilities/math-utils"
 
 export const UNDEF_RESULT = ""
@@ -33,3 +34,23 @@ equal = (a: any, b: any): boolean => {
 export const evaluateNode = (node: MathNode, scope?: FormulaMathJsScope) => {
   return node.compile().evaluate(scope)
 }
+
+// This function is used to get the root scope (an instance of our FormulaMathJSScope) within custom MathJS functions.
+// When the formula expression is executed, the initially passed scope can be wrapped in MathJS's PartitionedMap, which
+// is used to store temporary values. This function retrieves the root scope from the PartitionedMap.
+// This approach has been recommended by a MathJS maintainer. For more details, see:
+// https://github.com/josdejong/mathjs/pull/3150#issuecomment-2248101774
+export const getRootScope = (currentScope: CurrentScope): FormulaMathJsScope => {
+  return isPartitionedMap(currentScope) ? getRootScope(currentScope.a) : currentScope
+}
+
+export const isPartitionedMap = (map: any): map is MathJSPartitionedMap => {
+  // We could also check isMap(map.a), but that makes tests and mocking more difficult.
+  // The current check is probably specific enough to be safe.
+  return map && typeof map.a === "object" && isMap(map.b)
+}
+
+export const isMap = (map: any) => {
+  const requiredMapMethods = ['get', 'set', 'has', 'delete', 'entries']
+  return map && typeof map === 'object' && requiredMapMethods.every(method => typeof map[method] === 'function')
+}
diff --git a/v3/src/models/formula/functions/lookup-functions.ts b/v3/src/models/formula/functions/lookup-functions.ts
index 7596689402..3bbdb5f5c9 100644
--- a/v3/src/models/formula/functions/lookup-functions.ts
+++ b/v3/src/models/formula/functions/lookup-functions.ts
@@ -1,7 +1,7 @@
 import { ConstantNode, MathNode } from "mathjs"
-import { DisplayNameMap, FValue, ILookupDependency, MathJSPartitionedMap } from "../formula-types"
+import { CurrentScope, DisplayNameMap, FValue, ILookupDependency } from "../formula-types"
 import { rmCanonicalPrefix } from "../utils/name-mapping-utils"
-import { UNDEF_RESULT, equal, evaluateNode } from "./function-utils"
+import { UNDEF_RESULT, equal, evaluateNode, getRootScope } from "./function-utils"
 import { isConstantStringNode } from "../utils/mathjs-utils"
 import { t } from "../../../utilities/translation/translate"
 import type { IDataSet } from "../../data/data-set"
@@ -31,8 +31,8 @@ export const lookupFunctions = {
         attrNameArg.value = displayNameMap.dataSet[dataSetName]?.attribute[attrName] || attrName
       }
     },
-    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
-      const scope = partitionedMap.a
+    evaluateRaw: (args: MathNode[], mathjs: any, currentScope: CurrentScope) => {
+      const scope = getRootScope(currentScope)
       const functionName = "lookupByIndex"
       const numOfReqArgs = lookupFunctions.lookupByIndex.numOfRequiredArguments
       if (args.length !== numOfReqArgs) {
@@ -91,8 +91,8 @@ export const lookupFunctions = {
         keyAttrNameArg.value = displayNameMap.dataSet[dataSetName]?.attribute[keyAttrName] || keyAttrName
       }
     },
-    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
-      const scope = partitionedMap.a
+    evaluateRaw: (args: MathNode[], mathjs: any, currentScope: CurrentScope) => {
+      const scope = getRootScope(currentScope)
       const functionName = "lookupByKey"
       const numOfReqArgs = lookupFunctions.lookupByKey.numOfRequiredArguments
       if (args.length !== numOfReqArgs) {
diff --git a/v3/src/models/formula/functions/math.test.ts b/v3/src/models/formula/functions/math.test.ts
index 1aa69da8f8..0b6a43d16f 100644
--- a/v3/src/models/formula/functions/math.test.ts
+++ b/v3/src/models/formula/functions/math.test.ts
@@ -16,12 +16,12 @@ describe("evaluateRawWithAggregateContext", () => {
         scope.aggregateContext = false
       }
     }
-    const partitionedMap = { a: scope } as any as MathJSPartitionedMap
+    const currentScope = { a: scope, b: new Map() } as any as MathJSPartitionedMap
     let result = false
     const mockFn = jest.fn(() => { result = scope.aggregateContext; return "" })
 
-    evaluateRawWithAggregateContext(mockFn)(args, mathjs, partitionedMap)
-    expect(mockFn).toHaveBeenCalledWith(args, mathjs, partitionedMap)
+    evaluateRawWithAggregateContext(mockFn)(args, mathjs, currentScope)
+    expect(mockFn).toHaveBeenCalledWith(args, mathjs, currentScope)
     expect(result).toBeTruthy()
   })
 })
@@ -33,12 +33,12 @@ describe("evaluateRawWithDefaultArg", () => {
     const scope = {
       defaultArgumentNode: { name: "default" }
     }
-    const partitionedMap = { a: scope } as any as MathJSPartitionedMap
+    const currentScope = { a: scope, b: new Map() } as any as MathJSPartitionedMap
     const mockFn = jest.fn((_args: MathNode[]) => (_args[0] as SymbolNode).name)
 
     const numOfReqArgs = 1
-    const res = evaluateRawWithDefaultArg(mockFn, numOfReqArgs)(args, mathjs, partitionedMap)
-    expect(mockFn).toHaveBeenCalledWith([scope.defaultArgumentNode], mathjs, partitionedMap)
+    const res = evaluateRawWithDefaultArg(mockFn, numOfReqArgs)(args, mathjs, currentScope)
+    expect(mockFn).toHaveBeenCalledWith([scope.defaultArgumentNode], mathjs, currentScope)
     expect(res).toEqual("default")
   })
 
@@ -48,12 +48,12 @@ describe("evaluateRawWithDefaultArg", () => {
     const scope = {
       defaultArgumentNode: { name: "default" }
     }
-    const partitionedMap = { a: scope } as any as MathJSPartitionedMap
+    const currentScope = { a: scope, b: new Map() } as any as MathJSPartitionedMap
     const mockFn = jest.fn((_args: MathNode[]) => (_args[0] as SymbolNode).name)
 
     const res =
-      evaluateRawWithDefaultArg(mockFn, 1)(args as any as MathNode[], mathjs, partitionedMap)
-    expect(mockFn).toHaveBeenCalledWith(args, mathjs, partitionedMap)
+      evaluateRawWithDefaultArg(mockFn, 1)(args as any as MathNode[], mathjs, currentScope)
+    expect(mockFn).toHaveBeenCalledWith(args, mathjs, currentScope)
     expect(res).toEqual("provided")
   })
 })
@@ -63,10 +63,10 @@ describe("evaluateToEvaluateRaw", () => {
     const args = [ parse("1"), parse("2") ]
     const mathjs = {}
     const scope = {}
-    const partitionedMap = { a: scope } as any as MathJSPartitionedMap
+    const currentScope = { a: scope, b: new Map() } as any as MathJSPartitionedMap
     const mockFn = jest.fn((a: FValueOrArray, b: FValueOrArray) => Number(a) + Number(b))
 
-    const res = evaluateToEvaluateRaw(mockFn)(args as any as MathNode[], mathjs, partitionedMap)
+    const res = evaluateToEvaluateRaw(mockFn)(args as any as MathNode[], mathjs, currentScope)
     expect(mockFn).toHaveBeenCalledWith(1, 2)
     expect(res).toEqual(3)
   })
diff --git a/v3/src/models/formula/functions/math.ts b/v3/src/models/formula/functions/math.ts
index 55b76965a4..d3917491c6 100644
--- a/v3/src/models/formula/functions/math.ts
+++ b/v3/src/models/formula/functions/math.ts
@@ -1,9 +1,9 @@
 import { create, all, MathNode } from 'mathjs'
 import {
-  CODAPMathjsFunctionRegistry, EvaluateFunc, EvaluateFuncWithAggregateContextSupport, EvaluateRawFunc, FValue,
-  FValueOrArray, MathJSPartitionedMap
+  CODAPMathjsFunctionRegistry, CurrentScope, EvaluateFunc, EvaluateFuncWithAggregateContextSupport, EvaluateRawFunc,
+  FValue, FValueOrArray
 } from '../formula-types'
-import { evaluateNode } from './function-utils'
+import { evaluateNode, getRootScope } from './function-utils'
 import { arithmeticFunctions } from './arithmetic-functions'
 import { dateFunctions } from './date-functions'
 import { stringFunctions } from './string-functions'
@@ -17,26 +17,26 @@ export const math = create(all)
 
 // Each aggregate function needs to be evaluated with `withAggregateContext` method.
 export const evaluateRawWithAggregateContext = (fn: EvaluateRawFunc): EvaluateRawFunc => {
-  return (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
-    const scope = partitionedMap.a
+  return (args: MathNode[], mathjs: any, currentScope: CurrentScope) => {
+    const scope = getRootScope(currentScope)
     // withAggregateContext returns result of the callback function
-    return scope.withAggregateContext(() => fn(args, mathjs, partitionedMap))
+    return scope.withAggregateContext(() => fn(args, mathjs, currentScope))
   }
 }
 
 export const evaluateRawWithDefaultArg = (fn: EvaluateRawFunc, numOfRequiredArguments: number): EvaluateRawFunc => {
-  return (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
-    const scope = partitionedMap.a
+  return (args: MathNode[], mathjs: any, currentScope: CurrentScope) => {
+    const scope = getRootScope(currentScope)
     if (scope.defaultArgumentNode && args.length < numOfRequiredArguments) {
-      return fn([...args, scope.defaultArgumentNode], mathjs, partitionedMap)
+      return fn([...args, scope.defaultArgumentNode], mathjs, currentScope)
     }
-    return fn(args, mathjs, partitionedMap)
+    return fn(args, mathjs, currentScope)
   }
 }
 
 export const evaluateToEvaluateRaw = (fn: EvaluateFuncWithAggregateContextSupport): EvaluateRawFunc => {
-  return (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
-    const scope = partitionedMap.a
+  return (args: MathNode[], mathjs: any, currentScope: CurrentScope) => {
+    const scope = getRootScope(currentScope)
     return fn(...(args.map(arg => evaluateNode(arg, scope))))
   }
 }
diff --git a/v3/src/models/formula/functions/other-functions.ts b/v3/src/models/formula/functions/other-functions.ts
index ad4fe2916b..cd0e5e77ea 100644
--- a/v3/src/models/formula/functions/other-functions.ts
+++ b/v3/src/models/formula/functions/other-functions.ts
@@ -51,4 +51,6 @@ export const otherFunctions = {
     }
   },
 
+
+
 }
diff --git a/v3/src/models/formula/functions/semi-aggregate-functions.ts b/v3/src/models/formula/functions/semi-aggregate-functions.ts
index 9a47c9e774..a37db4dde4 100644
--- a/v3/src/models/formula/functions/semi-aggregate-functions.ts
+++ b/v3/src/models/formula/functions/semi-aggregate-functions.ts
@@ -1,19 +1,19 @@
 import { MathNode } from "mathjs"
-import { FValue, MathJSPartitionedMap } from "../formula-types"
-import { UNDEF_RESULT, evaluateNode, isValueTruthy } from "./function-utils"
+import { CurrentScope, FValue } from "../formula-types"
+import { UNDEF_RESULT, evaluateNode, getRootScope, isValueTruthy } from "./function-utils"
 
 export const semiAggregateFunctions = {
   next: {
     numOfRequiredArguments: 1,
     // expression and filter are evaluated as aggregate symbols, defaultValue is not - it depends on case index
     isSemiAggregate: [true, false, true],
-    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
+    evaluateRaw: (args: MathNode[], mathjs: any, currentScope: CurrentScope) => {
       interface ICachedData {
         result?: FValue
         resultCasePointer: number
       }
 
-      const scope = partitionedMap.a
+      const scope = getRootScope(currentScope)
       const caseGroupId = scope.getCaseGroupId()
       const cacheKey = `next(${args.toString()})-${caseGroupId}`
       const [ expression, defaultValue, filter ] = args
@@ -75,8 +75,8 @@ export const semiAggregateFunctions = {
     selfReferenceAllowed: true,
     // expression and filter are evaluated as aggregate symbols, defaultValue is not - it depends on case index
     isSemiAggregate: [true, false, true],
-    evaluateRaw: (args: MathNode[], mathjs: any, partitionedMap: MathJSPartitionedMap) => {
-      const scope = partitionedMap.a
+    evaluateRaw: (args: MathNode[], mathjs: any, currentScope: CurrentScope) => {
+      const scope = getRootScope(currentScope)
       const [ expression, defaultValue, filter ] = args
 
       const caseGroupId = scope.getCaseGroupId()

From daa947db3b103d254d889e39abb7834a879c6cda Mon Sep 17 00:00:00 2001
From: Piotr Janik <janikpiotrek@gmail.com>
Date: Thu, 25 Jul 2024 12:28:24 +0200
Subject: [PATCH 4/5] Update v3/src/models/formula/functions/function-utils.ts

Co-authored-by: Kirk Swenson <kswenson@users.noreply.github.com>
---
 v3/src/models/formula/functions/function-utils.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/v3/src/models/formula/functions/function-utils.ts b/v3/src/models/formula/functions/function-utils.ts
index 9bba84d34e..cd72b438e8 100644
--- a/v3/src/models/formula/functions/function-utils.ts
+++ b/v3/src/models/formula/functions/function-utils.ts
@@ -38,7 +38,7 @@ export const evaluateNode = (node: MathNode, scope?: FormulaMathJsScope) => {
 // This function is used to get the root scope (an instance of our FormulaMathJSScope) within custom MathJS functions.
 // When the formula expression is executed, the initially passed scope can be wrapped in MathJS's PartitionedMap, which
 // is used to store temporary values. This function retrieves the root scope from the PartitionedMap.
-// This approach has been recommended by a MathJS maintainer. For more details, see:
+// This approach has been recommended by the MathJS maintainer. For more details, see:
 // https://github.com/josdejong/mathjs/pull/3150#issuecomment-2248101774
 export const getRootScope = (currentScope: CurrentScope): FormulaMathJsScope => {
   return isPartitionedMap(currentScope) ? getRootScope(currentScope.a) : currentScope

From ec525490d25279cdb0c79c3f021d3b1de08752dd Mon Sep 17 00:00:00 2001
From: Piotr Janik <janikpiotrek@gmail.com>
Date: Thu, 25 Jul 2024 12:28:32 +0200
Subject: [PATCH 5/5] Update v3/src/models/formula/functions/other-functions.ts

Co-authored-by: Kirk Swenson <kswenson@users.noreply.github.com>
---
 v3/src/models/formula/functions/other-functions.ts | 2 --
 1 file changed, 2 deletions(-)

diff --git a/v3/src/models/formula/functions/other-functions.ts b/v3/src/models/formula/functions/other-functions.ts
index cd0e5e77ea..ad4fe2916b 100644
--- a/v3/src/models/formula/functions/other-functions.ts
+++ b/v3/src/models/formula/functions/other-functions.ts
@@ -51,6 +51,4 @@ export const otherFunctions = {
     }
   },
 
-
-
 }