diff --git a/src/TwigExtension/FlagCount.php b/src/TwigExtension/FlagCount.php
index e76b896..cada930 100644
--- a/src/TwigExtension/FlagCount.php
+++ b/src/TwigExtension/FlagCount.php
@@ -3,20 +3,21 @@
 namespace Drupal\flag\TwigExtension;
 
 use Drupal\Core\Entity\EntityInterface;
-use Drupal\Core\Template\TwigExtension;
 use Drupal\flag\FlagInterface;
+use Twig\TwigFunction;
+use Twig\Extension\AbstractExtension;
 
 /**
  * Provides a Twig extension to get the flag count given a flag and flaggable.
  */
-class FlagCount extends TwigExtension {
+class FlagCount extends AbstractExtension {
 
   /**
    * Generates a list of all Twig functions that this extension defines.
    */
   public function getFunctions() {
     return [
-      new \Twig_SimpleFunction('flagcount', [$this, 'count'], ['is_safe' => ['html']]),
+      new TwigFunction('flagcount', [$this, 'count'], ['is_safe' => ['html']]),
     ];
   }