Fix search function for new toc (#6065)

* Fix search function

* Small fix
diff --git a/docs/_static/js/sidebar.js b/docs/_static/js/sidebar.js
index 02c51ca..170a58d 100644
--- a/docs/_static/js/sidebar.js
+++ b/docs/_static/js/sidebar.js
@@ -41,7 +41,6 @@
                         render_left_helper($($.parseHTML(data)).find('#table-of-contents > div > ul'), currentText);
                     }
                     else {
-                        console.log(currentText);
                         render_left_helper($($.parseHTML(data)).find('.leftsidebar > .sphinxsidebarwrapper > ul.current > li.current > ul'), currentText);
                     }
                     keepExpand();
@@ -84,7 +83,6 @@
     var links = $('.rightsidebar a');
     for(var i = 1; i < links.length; ++i) {
         var divID = links.eq(i).attr('href');
-        console.log($(divID).offset().top - $(window).scrollTop());
         if($(divID).offset().top - $(window).scrollTop() > navbarHeight) {
             $('.rightsidebar a').css('color', 'black');
             links.eq(i - 1).css('color', '#337ab7');
@@ -146,12 +144,15 @@
 /*Keep toc expansion while redirecting*/
 function keepExpand() {
     var url = window.location.href, currentEntry;
-    $('.sphinxsidebar li').each(function () {
-        if(url.endsWith($(this).find('a').first().attr('href'))) {
-            currentEntry = $(this);
-            return false;
+    var entryList = isAPI ? $('.leftsidebar li') : $('.sphinxsidebar li');
+    for(var i = entryList.length - 1; i >= 0; --i) {
+        var entryURL = entryList.eq(i).find('a').first().attr('href');
+        if(entryURL != '#' && url.indexOf(entryURL) != -1) {
+            currentEntry = entryList.eq(i);
+            break;
         }
-    });
+    }
+    
     if(isAPI) {
         var rootEntry = currentEntry;
         if(rootEntry.parent().parent().is('li')) rootEntry = rootEntry.parent().parent();
diff --git a/docs/_static/searchtools_custom.js b/docs/_static/searchtools_custom.js
index 71ebed0..b2a17af 100644
--- a/docs/_static/searchtools_custom.js
+++ b/docs/_static/searchtools_custom.js
@@ -432,12 +432,14 @@
                 dirname = '';
               }
               listItem.append($('<a/>').attr('href',
-                DOCUMENTATION_OPTIONS.URL_ROOT + dirname).html(item[1]));
+                DOCUMENTATION_OPTIONS.URL_ROOT + dirname +
+                highlightstring + item[2]).html(item[1]));
             } else {
                 // normal html builders
                 var baseURL = 'http://mxnet.io/';
                 listItem.append($('<a/>').attr('href',
-                baseURL + item[0] + DOCUMENTATION_OPTIONS.FILE_SUFFIX).html(item[1]));
+                baseURL + item[0] + DOCUMENTATION_OPTIONS.FILE_SUFFIX +
+                highlightstring + item[2]).html(item[1]));
             }
             var html = listItem.children().html();
             listItem.children().html(html.replace(RegExp(query, 'gi'), '<strong>$&</strong>'));
@@ -475,11 +477,13 @@
             dirname = '';
           }
           listItem.append($('<a/>').attr('href',
-            DOCUMENTATION_OPTIONS.URL_ROOT + dirname).html(item[1]));
+            DOCUMENTATION_OPTIONS.URL_ROOT + dirname +
+            highlightstring + item[2]).html(item[1]));
         } else {
           // normal html builders
           listItem.append($('<a/>').attr('href',
-            item[0] + DOCUMENTATION_OPTIONS.FILE_SUFFIX).html(item[1]));
+            item[0] + DOCUMENTATION_OPTIONS.FILE_SUFFIX +
+            highlightstring + item[2]).html(item[1]));
         }
         if (item[3]) {
           listItem.append($('<span> (' + item[3] + ')</span>'));