1.before.cs

    private static void CheckUnsafeContext([NotNull] IAwaitExpression awaitExpression, [NotNull] IHighlightingConsumer consumer)
    {
      for (var localFunctionDeclaration = awaitExpression.GetContainingNode<ILocalFunctionDeclaration>();
        localFunctionDeclaration != null;
        localFunctionDeclaration = localFunctionDeclaration.GetContainingNode<ILocalFunctionDeclaration>())
      {
        if (localFunctionDeclaration.IsUnsafe)
        {
          consumer.AddHighlighting(new CannotAwaitInUnsafeContextError(awaitExpression));
          return;
        }
      }

      if (awaitExpression.GetContainingNode<IUnsafeCodeUnsafeStatement>() != null)
      {
        consumer.AddHighlighting(new CannotAwaitInUnsafeContextError(awaitExpression));
        return;
      }

      var memberDeclaration = awaitExpression.GetContainingNode<ICSharpTypeMemberDeclaration>(returnThis: true);
      if (memberDeclaration == null) return;

      if (memberDeclaration.IsUnsafe)
      {
        consumer.AddHighlighting(new CannotAwaitInUnsafeContextError(awaitExpression));
        return;
      }

      for (var typeDeclaration = memberDeclaration.GetContainingTypeDeclaration();
           typeDeclaration != null;
           typeDeclaration = typeDeclaration.GetContainingTypeDeclaration())
      {
        if (typeDeclaration.IsUnsafe)
        {
          consumer.AddHighlighting(new CannotAwaitInUnsafeContextError(awaitExpression));
          return;
        }
      }
    }

2.after.cs

    private static void CheckUnsafeContext([NotNull] IAwaitExpression awaitExpression, [NotNull] IHighlightingConsumer consumer)
    {
      foreach (var containingNode in awaitExpression.ContainingNodes())
      {
        switch (containingNode)
        {
          case ILocalFunctionDeclaration localFunctionDeclaration when localFunctionDeclaration.IsUnsafe:
          case IUnsafeCodeUnsafeStatement _:
          case ICSharpTypeMemberDeclaration memberDeclaration when memberDeclaration.IsUnsafe:
          case ICSharpTypeDeclaration typeDeclaration when typeDeclaration.IsUnsafe:
          {
            consumer.AddHighlighting(new CannotAwaitInUnsafeContextError(awaitExpression));
            return;
          }
        }
      }
    }