diff --git a/src/TaxonomyTermTree.php b/src/TaxonomyTermTree.php
index a4993d6b..3e211d47 100644
--- a/src/TaxonomyTermTree.php
+++ b/src/TaxonomyTermTree.php
@@ -2,73 +2,90 @@
 
 namespace Drupal\taxonomy_tree;
 
+use Drupal\Component\Plugin\Exception\InvalidPluginDefinitionException;
+use Drupal\Component\Plugin\Exception\PluginNotFoundException;
 use Drupal\Core\Entity\EntityTypeManager;
 
 /**
  * Loads taxonomy terms in a tree
  */
-class TaxonomyTermTree {
+class TaxonomyTermTree
+{
 
-  /**
-   * @var \Drupal\Core\Entity\EntityTypeManager
-   */
-  protected $entityTypeManager;
+    /**
+     * @var EntityTypeManager
+     */
+    protected $entityTypeManager;
 
-  /**
-   * TaxonomyTermTree constructor.
-   *
-   * @param \Drupal\Core\Entity\EntityTypeManager $entityTypeManager
-   */
-  public function __construct(EntityTypeManager $entityTypeManager) {
-    $this->entityTypeManager = $entityTypeManager;
-  }
-
-  /**
-   * Loads the tree of a vocabulary.
-   *
-   * @param string $vocabulary
-   *   Machine name
-   *
-   * @return array
-   */
-  public function load($vocabulary) {
-    $terms = $this->entityTypeManager->getStorage('taxonomy_term')->loadTree($vocabulary);
-    $tree = [];
-    foreach ($terms as $tree_object) {
-      $this->buildTree($tree, $tree_object, $vocabulary);
+    /**
+     * TaxonomyTermTree constructor.
+     *
+     * @param EntityTypeManager $entityTypeManager
+     */
+    public function __construct(EntityTypeManager $entityTypeManager)
+    {
+        $this->entityTypeManager = $entityTypeManager;
     }
 
-    return $tree;
-  }
+    /**
+     * Loads the tree of a vocabulary.
+     *
+     * @param $vocabulary
+     * @return array
+     * @throws InvalidPluginDefinitionException
+     * @throws PluginNotFoundException
+     */
+    public function load($vocabulary): array
+    {
+        $terms = $this->entityTypeManager->getStorage('taxonomy_term')->loadTree($vocabulary);
+        $tree = [];
+        foreach ($terms as $tree_object) {
+            $this->buildTree($tree, $tree_object, $vocabulary);
+        }
 
-  /**
-   * Populates a tree array given a taxonomy term tree object.
-   *
-   * @param $tree
-   * @param $object
-   * @param $vocabulary
-   */
-  protected function buildTree(&$tree, $object, $vocabulary) {
-    if ($object->depth != 0) {
-      return;
+        return $tree;
     }
-    $tree[$object->tid] = $object;
-    $tree[$object->tid]->children = [];
-    $object_children = &$tree[$object->tid]->children;
 
-    $children = $this->entityTypeManager->getStorage('taxonomy_term')->loadChildren($object->tid);
-    if (!$children) {
-      return;
-    }
+    /**
+     * Populates a tree array given a taxonomy term tree object.
+     *
+     * @param $tree
+     * @param $object
+     * @param $vocabulary
+     * @throws InvalidPluginDefinitionException
+     * @throws PluginNotFoundException
+     */
+    protected function buildTree(&$tree, $object, $vocabulary): void
+    {
+        if ($object->depth != 0) {
+            return;
+        }
+        $key = 'tid_' . $object->tid;
 
-    $child_tree_objects = $this->entityTypeManager->getStorage('taxonomy_term')->loadTree($vocabulary, $object->tid);
+        $tree[$key] = $object;
+        $tree[$key]->children = [];
+        $object_children = &$tree[$key]->children;
 
-    foreach ($children as $child) {
-      foreach ($child_tree_objects as $child_tree_object) {
-        if ($child_tree_object->tid == $child->id()) {
-         $this->buildTree($object_children, $child_tree_object, $vocabulary);
+        $children = $this->entityTypeManager->getStorage('taxonomy_term')->loadChildren($object->tid);
+        if (!$children) {
+            return;
         }
-      }
+
+        $child_tree_objects = $this->entityTypeManager->getStorage('taxonomy_term')->loadTree($vocabulary, $object->tid);
+
+        foreach ($children as $child) {
+            foreach ($child_tree_objects as $child_tree_object) {
+                if ($child_tree_object->tid == $child->id()) {
+                    $this->buildTree($object_children, $child_tree_object, $vocabulary);
+                }
+            }
+        }
+
+        uasort($tree, function ($a, $b) {
+            return $a->weight <=> $b->weight;
+        });
+        uasort($object_children, function ($a, $b) {
+            return $a->weight <=> $b->weight;
+        });
     }
-  }
-}
\ No newline at end of file
+}