TokenLayoutManager.swift

extension NSAttributedString.Key {
    static let token = NSAttributedString.Key("Token")
}

final class TokenLayoutManager: NSLayoutManager {
    var textContainerOriginOffset: CGSize = .zero
    
    override func drawGlyphs(forGlyphRange glyphsToShow: NSRange, at origin: CGPoint) {
        let characterRange = self.characterRange(forGlyphRange: glyphsToShow, actualGlyphRange: nil)
        textStorage?.enumerateAttribute(.token, in: characterRange, options: .longestEffectiveRangeNotRequired, using: { (value, subrange, _) in
            guard let token = value as? String, !token.isEmpty else { return }
            let tokenGlypeRange = glyphRange(forCharacterRange: subrange, actualCharacterRange: nil)
            drawToken(forGlyphRange: tokenGlypeRange)
        })
        super.drawGlyphs(forGlyphRange: glyphsToShow, at: origin)
    }
    
    private func drawToken(forGlyphRange tokenGlypeRange: NSRange) {
        guard let textContainer = textContainer(forGlyphAt: tokenGlypeRange.location, effectiveRange: nil) else { return }
        let withinRange = NSRange(location: NSNotFound, length: 0)
        enumerateEnclosingRects(forGlyphRange: tokenGlypeRange, withinSelectedGlyphRange: withinRange, in: textContainer) { (rect, _) in
            let tokenRect = rect.offsetBy(dx: self.textContainerOriginOffset.width, dy: self.textContainerOriginOffset.height)
            UIColor(hue: 175.0/360.0, saturation: 0.24, brightness: 0.88, alpha: 1).setFill()
            UIBezierPath(roundedRect: tokenRect, cornerRadius: 4).fill()
        }
    }
}