update elements and fix some bugs

lib/src/prism/plugins/keep-markup/prism-keep-markup.js

+(function () {
+
+        if (typeof self === 'undefined' || !self.Prism || !self.document || !document.createRange) {
+                return;
+        }
+
+        Prism.hooks.add('before-highlight', function (env) {
+                var firstWhiteSpaces = false;
+                var pos = 0;
+                var data = [];
+                var f = function (elt, baseNode) {
+                        var o = {};
+                        if (!baseNode) {
+                                // Clone the original tag to keep all attributes
+                                o.clone = elt.cloneNode(false);
+                                o.posOpen = pos;
+                                data.push(o);
+                        }
+                        for (var i = 0, l = elt.childNodes.length; i < l; i++) {
+                                var child = elt.childNodes[i];
+                                if (child.nodeType === 1) { // element
+                                        f(child);
+                                } else if(child.nodeType === 3) { // text
+                                        if(!firstWhiteSpaces) {
+                                                // We need to ignore the first white spaces in the code block
+                                                child.data = child.data.replace(/^(?:\r?\n|\r)/, '');
+                                                firstWhiteSpaces = true;
+                                        }
+                                        pos += child.data.length;
+                                }
+                        }
+                        if (!baseNode) {
+                                o.posClose = pos;
+                        }
+                };
+                f(env.element, true);
+
+                if (data && 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 clone
+                                                var range = document.createRange();
+                                                range.setStart(nodeState.nodeStart, nodeState.nodeStartPos);
+                                                range.setEnd(nodeState.nodeEnd, nodeState.nodeEndPos);
+                                                nodeState.node.clone.appendChild(range.extractContents());
+                                                range.insertNode(nodeState.node.clone);
+                                                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
+                                });
+                        });
+                }
+        });
+}());