diff --git a/src/TreeBuilder.php b/src/TreeBuilder.php
index 59b32ae..3f68c63 100644
--- a/src/TreeBuilder.php
+++ b/src/TreeBuilder.php
@@ -8,13 +8,25 @@
 namespace Drupal\token;
 
 use Drupal\Core\Cache\Cache;
-use Symfony\Component\DependencyInjection\ContainerAwareInterface;
-use Symfony\Component\DependencyInjection\ContainerAwareTrait;
-use Symfony\Component\DependencyInjection\ContainerInterface;
+use Drupal\Core\Cache\CacheBackendInterface;
+use Drupal\Core\Language\LanguageManagerInterface;
 
-class TreeBuilder implements TreeBuilderInterface, ContainerAwareInterface {
+class TreeBuilder implements TreeBuilderInterface {
 
-  use ContainerAwareTrait;
+  /**
+   * @var \Drupal\token\Token
+   */
+  protected $tokenService;
+
+  /**
+   * @var \Drupal\Core\Language\LanguageManagerInterface
+   */
+  protected $languageManager;
+
+  /**
+   * @var \Drupal\Core\Cache\CacheBackendInterface
+   */
+  protected $cacheBackend;
 
   /**
    * Cache already built trees.
@@ -23,16 +35,16 @@ class TreeBuilder implements TreeBuilderInterface, ContainerAwareInterface {
    */
   protected $builtTrees;
 
-  public function __construct(ContainerInterface $container) {
-    $this->container = $container;
+  public function __construct(TokenInterface $token_service, CacheBackendInterface $cache_backend, LanguageManagerInterface $language_manager) {
+    $this->tokenService = $token_service;
+    $this->cacheBackend = $cache_backend;
+    $this->languageManager = $language_manager;
   }
 
   /**
    * {@inheritdoc}
    */
   public function buildTree($token_type, array $options = []) {
-    $trees = &$this->builtTrees;
-
     $options += [
       'restricted' => FALSE,
       'depth' => 4,
@@ -49,19 +61,19 @@ class TreeBuilder implements TreeBuilderInterface, ContainerAwareInterface {
       $token_type = $entity_token_type;
     }
 
-    $langcode = $this->container->get('language_manager')->getCurrentLanguage()->getId();
+    $langcode = $this->languageManager->getCurrentLanguage()->getId();
     $tree_cid = "token_tree:{$token_type}:{$langcode}:{$options['depth']}";
 
     // If we do not have this base tree in the static cache, check the cache
     // otherwise generate and store it in the cache.
     if (!isset($this->builtTrees[$tree_cid])) {
-      if ($cache = $this->container->get('cache.data')->get($tree_cid)) {
+      if ($cache = $this->cacheBackend->get($tree_cid)) {
         $this->builtTrees[$tree_cid] = $cache->data;
       }
       else {
         $options['parents'] = [];
         $this->builtTrees[$tree_cid] = $this->getTokenData($token_type, $options);
-        $this->container->get('cache.data')->set($tree_cid, $this->builtTrees[$tree_cid], Cache::PERMANENT, [Token::TOKEN_INFO_CACHE_TAG]);
+        $this->cacheBackend->set($tree_cid, $this->builtTrees[$tree_cid], Cache::PERMANENT, [Token::TOKEN_INFO_CACHE_TAG]);
       }
     }
 
@@ -84,7 +96,7 @@ class TreeBuilder implements TreeBuilderInterface, ContainerAwareInterface {
         }
       }
       if (!empty($token_values)) {
-        $token_values = $this->container->get('token')->generate($token_type, $token_values, $options['data']);
+        $token_values = $this->tokenService->generate($token_type, $token_values, $options['data']);
         foreach ($token_values as $token => $replacement) {
           $tree[$token]['value'] = $replacement;
         }
@@ -102,7 +114,7 @@ class TreeBuilder implements TreeBuilderInterface, ContainerAwareInterface {
     foreach ($tree as $token => $token_info) {
       $result[$token] = $token_info;
       if (isset($token_info['children']) && is_array($token_info['children'])) {
-        $result += static::flattenTree($token_info['children']);
+        $result += $this->flattenTree($token_info['children']);
       }
     }
     return $result;
@@ -127,7 +139,7 @@ class TreeBuilder implements TreeBuilderInterface, ContainerAwareInterface {
       'parents' => [],
     ];
 
-    $info = $this->container->get('token')->getInfo();
+    $info = $this->tokenService->getInfo();
     if ($options['depth'] <= 0 || !isset($info['types'][$token_type]) || !isset($info['tokens'][$token_type])) {
       return [];
     }
@@ -170,7 +182,7 @@ class TreeBuilder implements TreeBuilderInterface, ContainerAwareInterface {
         $child_options = $options;
         $child_options['depth']--;
         $child_options['parents'] = $token_parents;
-        $tree[$raw_token]['children'] = static::getTokenData($token_info['type'], $child_options);
+        $tree[$raw_token]['children'] = $this->getTokenData($token_info['type'], $child_options);
       }
     }
 
diff --git a/token.module b/token.module
index f5cf49c..5c079e0 100644
--- a/token.module
+++ b/token.module
@@ -633,15 +633,6 @@ function token_build_tree($token_type, array $options = array()) {
 }
 
 /**
- * Flatten a token tree.
- *
- * @deprecated
- */
-function token_flatten_tree($tree) {
-  return \Drupal::service('token.tree_builder')->flattenTree($tree);
-}
-
-/**
  * Get a translated book menu link by its mlid, without access checking.
  *
  * This function is a copy of book_link_load() but with its own cache and a
diff --git a/token.services.yml b/token.services.yml
index c402c2f..a905b9b 100644
--- a/token.services.yml
+++ b/token.services.yml
@@ -1,4 +1,4 @@
 services:
   token.tree_builder:
     class: Drupal\token\TreeBuilder
-    arguments: ['@service_container']
+    arguments: ['@token', '@cache.data', '@language_manager']