prism-keep-markup.js

(function () {

	if (typeof Prism === 'undefined' || typeof document === 'undefined' || !document.createRange) {
		return;
	}

	Prism.plugins.KeepMarkup = true;

	Prism.hooks.add('before-highlight', function (env) {
		if (!env.element.children.length) {
			return;
		}

		if (!Prism.util.isActive(env.element, 'keep-markup', true)) {
			return;
		}

		var dropTokens = Prism.util.isActive(env.element, 'drop-tokens', false);
		/**
		 * Returns whether the given element should be kept.
		 *
		 * @param {HTMLElement} element
		 * @returns {boolean}
		 */
		function shouldKeep(element) {
			if (dropTokens && element.nodeName.toLowerCase() === 'span' && element.classList.contains('token')) {
				return false;
			}
			return true;
		}

		var pos = 0;
		var data = [];
		function processElement(element) {
			if (!shouldKeep(element)) {
				// don't keep this element and just process its children
				processChildren(element);
				return;
			}

			var o = {
				// Store original element so we can restore it after highlighting
				element: element,
				posOpen: pos
			};
			data.push(o);

			processChildren(element);

			o.posClose = pos;
		}
		function processChildren(element) {
			for (var i = 0, l = element.childNodes.length; i < l; i++) {
				var child = element.childNodes[i];
				if (child.nodeType === 1) { // element
					processElement(child);
				} else if (child.nodeType === 3) { // text
					pos += child.data.length;
				}
			}
		}
		processChildren(env.element);

		if (data.length) {
			// data is an array of all existing tags
			env.keepMarkup = data;
		}
	});

	Prism.hooks.add('after-highlight', function (env) {
		if (env.keepMarkup && env.keepMarkup.length) {

			var walk = function (elt, nodeState) {
				for (var i = 0, l = elt.childNodes.length; i < l; i++) {

					var child = elt.childNodes[i];

					if (child.nodeType === 1) { // element
						if (!walk(child, nodeState)) {
							return false;
						}

					} else if (child.nodeType === 3) { // text
						if (!nodeState.nodeStart && nodeState.pos + child.data.length > nodeState.node.posOpen) {
							// We found the start position
							nodeState.nodeStart = child;
							nodeState.nodeStartPos = nodeState.node.posOpen - nodeState.pos;
						}
						if (nodeState.nodeStart && nodeState.pos + child.data.length >= nodeState.node.posClose) {
							// We found the end position
							nodeState.nodeEnd = child;
							nodeState.nodeEndPos = nodeState.node.posClose - nodeState.pos;
						}

						nodeState.pos += child.data.length;
					}

					if (nodeState.nodeStart && nodeState.nodeEnd) {
						// Select the range and wrap it with the element
						var range = document.createRange();
						range.setStart(nodeState.nodeStart, nodeState.nodeStartPos);
						range.setEnd(nodeState.nodeEnd, nodeState.nodeEndPos);
						nodeState.node.element.innerHTML = '';
						nodeState.node.element.appendChild(range.extractContents());
						range.insertNode(nodeState.node.element);
						range.detach();

						// Process is over
						return false;
					}
				}
				return true;
			};

			// For each tag, we walk the DOM to reinsert it
			env.keepMarkup.forEach(function (node) {
				walk(env.element, {
					node: node,
					pos: 0
				});
			});
			// Store new highlightedCode for later hooks calls
			env.highlightedCode = env.element.innerHTML;
		}
	});
}());