From e36bd98723d97eb9fd123527bd94ce50879c1526 Mon Sep 17 00:00:00 2001
From: DMartens <DMartens@users.noreply.github.com>
Date: Thu, 16 Apr 2020 00:33:30 +0200
Subject: [PATCH] Add updateNodeSourceCodeLocation to the tree adapter (#315)

* Add updateNodeSourceCodeLocation to the tree adapter

* Add updateNodeSourceCodeLocation to the tree adapter documentation

* Renamed updateNodeSourceCodeLocation parameter location to endLocation

* Use updateNodeSourceCodeLocation to update the end location of tokens

* Update tree adapter updateNodeSourceCodeLocation documentation to reflect updating the source code location of nodes not just elements

* Add test for updating node source code location
---
 .../lib/index.js                              |  4 ++
 .../docs/source-code-location/end-location.md | 48 ++++++++++++++++
 .../parse5/docs/tree-adapter/interface.md     | 19 ++++++-
 .../extensions/location-info/parser-mixin.js  | 23 ++++----
 packages/parse5/lib/tree-adapters/default.js  |  4 ++
 .../parse5/test/location-info-parser.test.js  | 55 ++++++++++++++++++-
 6 files changed, 139 insertions(+), 14 deletions(-)
 create mode 100644 packages/parse5/docs/source-code-location/end-location.md

diff --git a/packages/parse5-htmlparser2-tree-adapter/lib/index.js b/packages/parse5-htmlparser2-tree-adapter/lib/index.js
index 824047383..58464d78a 100644
--- a/packages/parse5-htmlparser2-tree-adapter/lib/index.js
+++ b/packages/parse5-htmlparser2-tree-adapter/lib/index.js
@@ -342,3 +342,7 @@ exports.setNodeSourceCodeLocation = function(node, location) {
 exports.getNodeSourceCodeLocation = function(node) {
     return node.sourceCodeLocation;
 };
+
+exports.updateNodeSourceCodeLocation = function(node, endLocation) {
+    node.sourceCodeLocation = Object.assign(node.sourceCodeLocation, endLocation);
+};
diff --git a/packages/parse5/docs/source-code-location/end-location.md b/packages/parse5/docs/source-code-location/end-location.md
new file mode 100644
index 000000000..8b3bd57db
--- /dev/null
+++ b/packages/parse5/docs/source-code-location/end-location.md
@@ -0,0 +1,48 @@
+# Interface: EndLocation
+
+### Properties
+
+* [endCol](#endcol)
+* [endOffset](#endoffset)
+* [endLine](#endline)
+* [endTag](#endtag)
+
+---
+
+## Properties
+
+<a id="endcol"></a>
+
+###  endCol
+
+**● endCol**: *`number`*
+
+One-based column index of the last character
+
+___
+<a id="endoffset"></a>
+
+###  endOffset
+
+**● endOffset**: *`number`*
+
+Zero-based last character index
+
+___
+<a id="endline"></a>
+
+###  endLine
+
+**● endLine**: *`number`*
+
+One-based line index of the last character
+
+___
+<a id="endtag"></a>
+
+###  endTag
+
+**● endTag**: *[Location](location.md)|undefined*
+
+Element's end tag location info.
+This property is undefined, if the element has no closing tag.
\ No newline at end of file
diff --git a/packages/parse5/docs/tree-adapter/interface.md b/packages/parse5/docs/tree-adapter/interface.md
index ac1e56542..f7ef11c47 100644
--- a/packages/parse5/docs/tree-adapter/interface.md
+++ b/packages/parse5/docs/tree-adapter/interface.md
@@ -38,7 +38,7 @@ Tree adapter is a set of utility functions that provides minimal required abstra
 * [setDocumentType](#setdocumenttype)
 * [setNodeSourceCodeLocation](#setnodesourcecodelocation)
 * [setTemplateContent](#settemplatecontent)
-
+* [updateNodeSourceCodeLocation](#updatenodesourcecodelocation)
 ---
 
 ## Methods
@@ -588,6 +588,21 @@ Sets the `<template>` element content element.
 | contentElement | DocumentFragment |  Content element. |
 
 **Returns:** `void`
-
 ___
+<a id="updatenodesourcecodelocation"></a>
+
+###  updateNodeSourceCodeLocation
+
+▸ **updateNodeSourceCodeLocation**(node: *Node*, endLocation: *[EndLocation](../source-code-location/end-location.md)*): `void`
+
+Updates the source code location of nodes.
+
+**Parameters:**
 
+| Param | Type | Description |
+| ------ | ------ | ------ |
+| node | Node |  Node. |
+| endLocation | [EndLocation](../source-code-location/end-location.md) |  Source code location information of the end of the node. |
+
+**Returns:** `void`
+___
diff --git a/packages/parse5/lib/extensions/location-info/parser-mixin.js b/packages/parse5/lib/extensions/location-info/parser-mixin.js
index fcf3a40af..e7d3e2da1 100644
--- a/packages/parse5/lib/extensions/location-info/parser-mixin.js
+++ b/packages/parse5/lib/extensions/location-info/parser-mixin.js
@@ -43,17 +43,19 @@ class LocationInfoParserMixin extends Mixin {
                 // NOTE: For cases like <p> <p> </p> - First 'p' closes without a closing
                 // tag and for cases like <td> <p> </td> - 'p' closes without a closing tag.
                 const isClosingEndTag = closingToken.type === Tokenizer.END_TAG_TOKEN && tn === closingToken.tagName;
-
+                const endLoc = {};
                 if (isClosingEndTag) {
-                    loc.endTag = Object.assign({}, ctLoc);
-                    loc.endLine = ctLoc.endLine;
-                    loc.endCol = ctLoc.endCol;
-                    loc.endOffset = ctLoc.endOffset;
+                    endLoc.endTag = Object.assign({}, ctLoc);
+                    endLoc.endLine = ctLoc.endLine;
+                    endLoc.endCol = ctLoc.endCol;
+                    endLoc.endOffset = ctLoc.endOffset;
                 } else {
-                    loc.endLine = ctLoc.startLine;
-                    loc.endCol = ctLoc.startCol;
-                    loc.endOffset = ctLoc.startOffset;
+                    endLoc.endLine = ctLoc.startLine;
+                    endLoc.endCol = ctLoc.startCol;
+                    endLoc.endOffset = ctLoc.startOffset;
                 }
+
+                this.treeAdapter.updateNodeSourceCodeLocation(element, endLoc);
             }
         }
     }
@@ -208,9 +210,8 @@ class LocationInfoParserMixin extends Mixin {
                 const tnLoc = this.treeAdapter.getNodeSourceCodeLocation(textNode);
 
                 if (tnLoc) {
-                    tnLoc.endLine = token.location.endLine;
-                    tnLoc.endCol = token.location.endCol;
-                    tnLoc.endOffset = token.location.endOffset;
+                    const { endLine, endCol, endOffset } = token.location;
+                    this.treeAdapter.updateNodeSourceCodeLocation(textNode, { endLine, endCol, endOffset });
                 } else {
                     this.treeAdapter.setNodeSourceCodeLocation(textNode, token.location);
                 }
diff --git a/packages/parse5/lib/tree-adapters/default.js b/packages/parse5/lib/tree-adapters/default.js
index 7935a0e0b..14d007a16 100644
--- a/packages/parse5/lib/tree-adapters/default.js
+++ b/packages/parse5/lib/tree-adapters/default.js
@@ -215,3 +215,7 @@ exports.setNodeSourceCodeLocation = function(node, location) {
 exports.getNodeSourceCodeLocation = function(node) {
     return node.sourceCodeLocation;
 };
+
+exports.updateNodeSourceCodeLocation = function(node, endLocation) {
+    node.sourceCodeLocation = Object.assign(node.sourceCodeLocation, endLocation);
+};
diff --git a/packages/parse5/test/location-info-parser.test.js b/packages/parse5/test/location-info-parser.test.js
index 22f474c86..4d4ce6879 100644
--- a/packages/parse5/test/location-info-parser.test.js
+++ b/packages/parse5/test/location-info-parser.test.js
@@ -7,7 +7,7 @@ const {
     assertStartTagLocation,
     assertNodeLocation
 } = require('../../../test/utils/generate-location-info-parser-tests');
-const { generateTestsForEachTreeAdapter } = require('../../../test/utils/common');
+const { generateTestsForEachTreeAdapter, treeAdapters } = require('../../../test/utils/common');
 
 generateLocationInfoParserTests(module.exports, 'Parser', (input, opts) => ({
     node: parse5.parse(input, opts)
@@ -127,3 +127,56 @@ generateTestsForEachTreeAdapter(module.exports, (_test, treeAdapter) => {
         assert.ok(!location.endTag);
     };
 });
+
+exports['Updating node source code location (GH-314)'] = function() {
+    const sourceCodeLocationSetter = {
+        setNodeSourceCodeLocation(node, location) {
+            if (location === null) {
+                node.sourceCodeLocation = null;
+            } else {
+                node.sourceCodeLocation = {
+                    start: {
+                        line: location.startLine,
+                        column: location.startCol,
+                        offset: location.startOffset
+                    },
+                    end: {
+                        line: location.endLine,
+                        column: location.endCol,
+                        offset: location.endOffset
+                    }
+                };
+            }
+        },
+        updateNodeSourceCodeLocation(node, endLocation) {
+            node.sourceCodeLocation = {
+                start: node.sourceCodeLocation.start,
+                end: {
+                    line: endLocation.endLine,
+                    column: endLocation.endCol,
+                    offset: endLocation.endOffset
+                }
+            };
+        }
+    };
+    const adapter = Object.assign(treeAdapters.default, sourceCodeLocationSetter);
+    const document = parse5.parse('<!doctype><body>Testing location</body>', { adapter, sourceCodeLocationInfo: true });
+    const [doctype, html] = document.childNodes;
+    const [head, body] = html.childNodes;
+    const [text] = body.childNodes;
+
+    assert.deepEqual(doctype.sourceCodeLocation, {
+        start: { line: 1, column: 1, offset: 0 },
+        end: { line: 1, column: 11, offset: 10 }
+    });
+    assert.strictEqual(html.sourceCodeLocation, null);
+    assert.strictEqual(head.sourceCodeLocation, null);
+    assert.deepEqual(body.sourceCodeLocation, {
+        start: { line: 1, column: 11, offset: 10 },
+        end: { line: 1, column: 40, offset: 39 }
+    });
+    assert.deepEqual(text.sourceCodeLocation, {
+        start: { line: 1, column: 17, offset: 16 },
+        end: { line: 1, column: 33, offset: 32 }
+    });
+};