From 34f9ac0035da178527414973fa26a176e0089104 Mon Sep 17 00:00:00 2001
From: kurkle <jukka.kurkela@gmail.com>
Date: Wed, 21 Jul 2021 09:04:14 +0300
Subject: [PATCH] Add data/dataset indices to line segment context

---
 docs/charts/line.md                |  3 +++
 src/controllers/controller.line.js |  1 +
 src/elements/element.line.js       |  1 +
 src/helpers/helpers.segment.js     | 22 +++++++++++++++-------
 types/index.esm.d.ts               |  5 ++++-
 5 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/docs/charts/line.md b/docs/charts/line.md
index bc5f6ceb78e..d872dd9d6b1 100644
--- a/docs/charts/line.md
+++ b/docs/charts/line.md
@@ -168,6 +168,9 @@ Context for the scriptable segment contains the following properties:
 * `type`: `'segment'`
 * `p0`: first point element
 * `p1`: second point element
+* `p0DataIndex`: index of first point in the data array
+* `p1DataIndex`: index of second point in the data array
+* `datasetIndex`: dataset index
 
 [Example usage](../samples/line/segments.md)
 
diff --git a/src/controllers/controller.line.js b/src/controllers/controller.line.js
index ab3c5051556..9b6a505fe0b 100644
--- a/src/controllers/controller.line.js
+++ b/src/controllers/controller.line.js
@@ -27,6 +27,7 @@ export default class LineController extends DatasetController {
     }
 
     // Update Line
+    line._datasetIndex = me.index;
     line._decimated = !!_dataset._decimated;
     line.points = points;
 
diff --git a/src/elements/element.line.js b/src/elements/element.line.js
index d2aff52c8bc..c09548c41ce 100644
--- a/src/elements/element.line.js
+++ b/src/elements/element.line.js
@@ -251,6 +251,7 @@ export default class LineElement extends Element {
     this._segments = undefined;
     this._decimated = false;
     this._pointsUpdated = false;
+    this._datasetIndex = undefined;
 
     if (cfg) {
       Object.assign(this, cfg);
diff --git a/src/helpers/helpers.segment.js b/src/helpers/helpers.segment.js
index 215c73207ba..dd9f4da9f2c 100644
--- a/src/helpers/helpers.segment.js
+++ b/src/helpers/helpers.segment.js
@@ -244,15 +244,14 @@ export function _computeSegments(line, segmentOptions) {
 
   const loop = !!line._loop;
   const {start, end} = findStartAndEnd(points, count, loop, spanGaps);
-  const baseStyle = readStyle(line.options);
 
   if (spanGaps === true) {
-    return splitByStyles([{start, end, loop}], points, baseStyle, segmentOptions);
+    return splitByStyles(line, [{start, end, loop}], points, segmentOptions);
   }
 
   const max = end < start ? end + count : end;
   const completeLoop = !!line._fullLoop && start === 0 && end === count - 1;
-  return splitByStyles(solidSegments(points, start, max, completeLoop), points, baseStyle, segmentOptions);
+  return splitByStyles(line, solidSegments(points, start, max, completeLoop), points, segmentOptions);
 }
 
 /**
@@ -261,20 +260,22 @@ export function _computeSegments(line, segmentOptions) {
  * @param {object} [segmentOptions]
  * @return {Segment[]}
  */
-function splitByStyles(segments, points, baseStyle, segmentOptions) {
+function splitByStyles(line, segments, points, segmentOptions) {
   if (!segmentOptions || !segmentOptions.setContext || !points) {
     return segments;
   }
-  return doSplitByStyles(segments, points, baseStyle, segmentOptions);
+  return doSplitByStyles(line, segments, points, segmentOptions);
 }
 
 /**
+ * @param {LineElement} line
  * @param {Segment[]} segments
  * @param {PointElement[]} points
  * @param {object} [segmentOptions]
  * @return {Segment[]}
  */
-function doSplitByStyles(segments, points, baseStyle, segmentOptions) {
+function doSplitByStyles(line, segments, points, segmentOptions) {
+  const baseStyle = readStyle(line.options);
   const count = points.length;
   const result = [];
   let start = segments[0].start;
@@ -285,7 +286,14 @@ function doSplitByStyles(segments, points, baseStyle, segmentOptions) {
     let style;
     for (i = start + 1; i <= segment.end; i++) {
       const pt = points[i % count];
-      style = readStyle(segmentOptions.setContext({type: 'segment', p0: prev, p1: pt}));
+      style = readStyle(segmentOptions.setContext({
+        type: 'segment',
+        p0: prev,
+        p1: pt,
+        p0DataIndex: (i - 1) % count,
+        p1DataIndex: i % count,
+        datasetIndex: line._datasetIndex
+      }));
       if (styleChanged(style, prevStyle)) {
         result.push({start: start, end: i - 1, loop: segment.loop, style: prevStyle});
         prevStyle = style;
diff --git a/types/index.esm.d.ts b/types/index.esm.d.ts
index 4b1e36e402f..ccd1838cdb9 100644
--- a/types/index.esm.d.ts
+++ b/types/index.esm.d.ts
@@ -28,7 +28,10 @@ export interface ScriptableContext<TType extends ChartType> {
 export interface ScriptableLineSegmentContext {
   type: 'segment',
   p0: PointElement,
-  p1: PointElement
+  p1: PointElement,
+  p0DataIndex: number,
+  p1DataIndex: number,
+  datasetIndex: number
 }
 
 export type Scriptable<T, TContext> = T | ((ctx: TContext, options: AnyObject) => T);