Regex.hs

{-|
Copyright        : (c) Galois, Inc. 2025
Maintainer       : Langston Barrett <langston@galois.com>
-}

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module Lang.Crucible.Debug.Regex
  ( type Regex
  , type Empty
  , type Lit
  , type (:|)
  , type Then
  , type Star
  , RegexRepr(..)
  , AcceptsEmpty
  , acceptsEmpty
  , DerivativeResult(..)
  , derivative
  , Match(..)
  , MatchError(..)
  , TokenParser(..)
  , match
  , nextLits
  , liftOrs
  , Sugar(..)
  , sugar
  , prettySugar
  ) where

import Data.Bifunctor (first)
import Data.Foldable qualified as Foldable
import Data.Kind (Type)
import Data.Parameterized.BoolRepr (BoolRepr (..), (%||), (%&&))
import Data.Parameterized.Classes (KnownRepr(knownRepr))
import Data.Parameterized.Some (Some (Some))
import Data.Parameterized.TraversableFC qualified as TFC
import Data.Sequence qualified as Seq
import Data.Sequence (Seq)
import Data.Type.Bool (type (||), type (&&))
import Data.Type.Equality (TestEquality (testEquality), type (:~:)(Refl))
import Prettyprinter qualified as PP

-- | Type-level only
data Regex a
  = TFail
  | TEmpty
  | TLit a
  | TOr (Regex a) (Regex a)
  | TThen (Regex a) (Regex a)
  | TStar (Regex a)

type Empty = 'TEmpty
type Lit = 'TLit
type a :| b = 'TOr a b
type Then = 'TThen
type Star = 'TStar

-- | Value-level representative of 'Regex'
--
-- The order of the type parameters is a bit arbitrary... This order gives
-- 'KnownRepr' and 'TestEquality' instances but requires flipping the type
-- parameters to get 'TraversableF'.
type RegexRepr :: (k -> Type) -> Regex k -> Type
data RegexRepr f r where
  Empty :: RegexRepr f 'TEmpty
  Fail :: RegexRepr f 'TFail
  Lit :: f t -> RegexRepr f ('TLit t)
  Or :: RegexRepr f l -> RegexRepr f r -> RegexRepr f ('TOr l r)
  Star :: RegexRepr f r -> RegexRepr f ('TStar r)
  Then :: RegexRepr f l -> RegexRepr f r -> RegexRepr f ('TThen l r)

instance TestEquality f => TestEquality (RegexRepr f) where
  testEquality rgx rgx' =
    case (rgx, rgx') of
      (Empty, Empty) -> Just Refl
      (Empty, _) -> Nothing
      (Fail, Fail) -> Just Refl
      (Fail, _) -> Nothing
      (Lit l, Lit l') ->
        case testEquality l l' of
          Just Refl -> Just Refl
          Nothing -> Nothing
      (Lit {}, _) -> Nothing
      (Or l r, Or l' r') ->
        case (testEquality l l', testEquality r r') of
          (Just Refl, Just Refl) -> Just Refl
          _ -> Nothing
      (Or {}, _) -> Nothing
      (Star r, Star r') ->
        case testEquality r r' of
          Just Refl -> Just Refl
          Nothing -> Nothing
      (Star {}, _) -> Nothing
      (Then l r, Then l' r') ->
        case (testEquality l l', testEquality r r') of
          (Just Refl, Just Refl) -> Just Refl
          _ -> Nothing
      (Then {}, _) -> Nothing

instance KnownRepr (RegexRepr f) 'TEmpty where
  knownRepr = Empty

instance KnownRepr (RegexRepr f) 'TFail where
  knownRepr = Fail

instance KnownRepr f t => KnownRepr (RegexRepr f) ('TLit t) where
  knownRepr = Lit knownRepr

instance
  ( KnownRepr (RegexRepr f) l
  , KnownRepr (RegexRepr f) r
  ) => KnownRepr (RegexRepr f) ('TOr l r) where
  knownRepr = Or knownRepr knownRepr

instance
  ( KnownRepr (RegexRepr f) l
  , KnownRepr (RegexRepr f) r
  ) => KnownRepr (RegexRepr f) ('TThen l r) where
  knownRepr = Then knownRepr knownRepr

instance KnownRepr (RegexRepr f) r => KnownRepr (RegexRepr f) ('TStar r) where
  knownRepr = Star knownRepr

instance TFC.FunctorFC RegexRepr where
  fmapFC f =
    \case
      Empty -> Empty
      Fail -> Fail
      Lit l -> Lit (f l)
      Or l r -> Or (TFC.fmapFC f l) (TFC.fmapFC f r)
      Star r -> Star (TFC.fmapFC f r)
      Then l r -> Then (TFC.fmapFC f l) (TFC.fmapFC f r)

instance TFC.FoldableFC RegexRepr where
  foldMapFC f =
    \case
      Empty -> mempty
      Fail -> mempty
      Lit l -> f l
      Or l r -> TFC.foldMapFC f l <> TFC.foldMapFC f r
      Star r -> TFC.foldMapFC f r
      Then l r -> TFC.foldMapFC f l <> TFC.foldMapFC f r

instance TFC.TraversableFC RegexRepr where
  traverseFC f =
    \case
      Empty -> pure Empty
      Fail -> pure Fail
      Lit l -> Lit <$> f l
      Or l r -> Or <$> TFC.traverseFC f l <*> TFC.traverseFC f r
      Star r -> Star <$> TFC.traverseFC f r
      Then l r -> Then <$> TFC.traverseFC f l <*> TFC.traverseFC f r

type AcceptsEmpty :: Regex k -> Bool
type family AcceptsEmpty r where
  AcceptsEmpty 'TFail = 'False
  AcceptsEmpty 'TEmpty  = 'True
  AcceptsEmpty ('TLit _) = 'False
  AcceptsEmpty ('TOr l r) = AcceptsEmpty l || AcceptsEmpty r
  AcceptsEmpty ('TThen l r) = AcceptsEmpty l && AcceptsEmpty r
  AcceptsEmpty ('TStar r) = 'True

acceptsEmpty :: RegexRepr f r -> BoolRepr (AcceptsEmpty r)
acceptsEmpty =
  \case
    Empty -> TrueRepr
    Fail -> FalseRepr
    Lit {} -> FalseRepr
    Or l r -> acceptsEmpty l %|| acceptsEmpty r
    Star _ -> TrueRepr
    Then l r -> acceptsEmpty l %&& acceptsEmpty r

-- | Type-level version of 'nu'.
--
-- See Wikipedia on "Brzozowski derivative"
type Nu :: Regex k -> Regex k
type family Nu r where
  Nu 'TFail = 'TFail
  Nu 'TEmpty  = 'TEmpty
  Nu ('TLit _) = 'TFail
  Nu ('TOr l r) = 'TOr (Nu l) (Nu r)
  Nu ('TThen l r) = 'TThen (Nu l) (Nu r)
  Nu ('TStar r) = 'TEmpty

-- | Auxiliary function for 'derivative'.
--
-- Value-level version of 'Nu'.
--
-- See Wikipedia on "Brzozowski derivative"
nu :: RegexRepr f r -> RegexRepr f (Nu r)
nu =
  \case
    Empty -> Empty
    Fail -> Fail
    Lit {} -> Fail
    Or l r -> Or (nu l) (nu r)
    Star _ -> Empty
    Then l r -> Then (nu l) (nu r)

-- | The result of 'derivative'
data DerivativeResult f g
  = DerivativeResult
    { -- | The remaining regex after matching that token
      derivativeRegex :: Some (RegexRepr f)
      -- | All the literals the token was matched at
    , derivativeMatched :: Seq (Some g)
    }

-- | See Wikipedia on "Brzozowski derivative"
derivative ::
  (forall t. f t -> TokenParser tok e g t) ->
  tok ->
  RegexRepr f r ->
  DerivativeResult f g
derivative getParser tok =
  let noMatch r = DerivativeResult r Seq.empty in
  let doFail = noMatch (Some Fail) in
  \case
    Empty -> doFail
    Fail -> doFail
    Lit f ->
      case runTokenParser (getParser f) tok of
        Left _ -> doFail
        Right m -> DerivativeResult (Some Empty) (Seq.singleton (Some m))
    Or l r ->
      case (derivative getParser tok l, derivative getParser tok r) of
        (DerivativeResult (Some Fail) _, r') -> r'
        (l', DerivativeResult (Some Fail) _) -> l'
        (DerivativeResult (Some l') ms, DerivativeResult (Some r') ms') ->
          DerivativeResult (Some (Or l' r')) (ms <> ms')
    Star f ->
      case derivative getParser tok f of
        DerivativeResult (Some f') ms ->
          DerivativeResult (Some (Then f' (Star f))) ms
    Then l r ->
      case (derivative getParser tok l, derivative getParser tok r) of
        (DerivativeResult (Some  l') ms, DerivativeResult (Some r') ms') ->
          DerivativeResult (Some (Or (Then l' r) (Then (nu l) r'))) (ms <> ms')

-- | The result of 'match', in case of success
type Match :: (k -> Type) -> Regex k -> Type
data Match k r where
  MEmpty :: Match f 'TEmpty
  MLit :: f t -> Match f ('TLit t)
  MLeft :: Match f l -> Match f ('TOr l r)
  MRight :: Match f r -> Match f ('TOr l r)
  MThen :: Match f l -> Match f r -> Match f ('TThen l r)
  MStar :: [Match f r] -> Match f ('TStar r)

-- | Failures that may arise during 'match'
data MatchError tok e
  = EFail
  | Eof
  | EOr (MatchError tok e) (MatchError tok e)
  | NotEmpty [tok]
  | Token e
  deriving Show

instance (PP.Pretty e, PP.Pretty tok) => PP.Pretty (MatchError tok e) where
  pretty =
    \case
      EFail -> "This regular expression never matches"
      Eof -> "Unexpected end of input"
      EOr l r ->
        PP.vcat
        [ "Both branches failed to match:"
        , "-" PP.<+> PP.align (PP.pretty l)
        , "-" PP.<+> PP.align (PP.pretty r)
        ]
      NotEmpty toks ->
        "Expected end of input, but found:" PP.<+> PP.hsep (map PP.pretty toks)
      Token e -> PP.pretty e

type TokenParser :: Type -> Type -> (k -> Type) -> k -> Type
newtype TokenParser tok e f t
  = TokenParser { runTokenParser :: tok -> Either e (f t) }

-- | Match a regular expression against a token stream
match ::
  -- | Regex
  RegexRepr (TokenParser tok e g) r ->
  -- | Token stream
  [tok] ->
  -- | Either a 'MatchError', or a 'Match' plus the unconsumed tokens
  Either (MatchError tok e) (Match g r, [tok])
match rgx toks =
  case (rgx, toks) of
    (Fail, _) -> Left EFail
    (Empty, []) -> Right (MEmpty, toks)
    (Empty, _) -> Left (NotEmpty toks)
    (Lit _, []) -> Left Eof
    (Lit ft, t : ts) -> do
      m <- first Token (runTokenParser ft t)
      pure (MLit m, ts)
    (Then l r, ts) -> do
      (ml, ts') <- match l ts
      (mr, ts'') <- match r ts'
      pure (MThen ml mr, ts'')
    (Or l r, ts) ->
      case match l ts of
        Right (m, ts') -> Right (MLeft m, ts')
        Left ef ->
          case match r ts of
            Right (m, ts') -> pure (MRight m, ts')
            Left eg -> Left (EOr ef eg)
    (Star f, ts) ->
      case match f ts of
        Left _ -> Right (MStar [], ts)
        Right (m, ts') ->
          first (\(MStar ms) -> MStar (m : ms)) <$> match (Star f) ts'

-- | List the literals that could be matched next
nextLits :: RegexRepr f r -> Seq (Some f)
nextLits =
  \case
    Empty -> Seq.empty
    Fail -> Seq.empty
    Lit x -> Seq.singleton (Some x)
    Or l r -> nextLits l <> nextLits r
    Star r -> nextLits r
    Then l r ->
      case acceptsEmpty l of
        TrueRepr -> nextLits l <> nextLits r
        FalseRepr -> nextLits l

-- | Syntactic sugar for displaying 'RegexRepr's, especially for quantification
type Sugar :: (k -> Type) -> Regex k -> Type
data Sugar f r where
  SEmpty :: Sugar f 'TEmpty
  SFail :: Sugar f 'TFail
  SLit :: f t -> Sugar f ('TLit t)
  SOptL :: Sugar f l -> Sugar f ('TOr l Empty)
  SOptR :: Sugar f r -> Sugar f ('TOr Empty r)
  SOr :: Sugar f l -> Sugar f r -> Sugar f ('TOr l r)
  SSomeL :: Sugar f l -> Sugar f ('TOr ('TStar r) r)
  SSomeR :: Sugar f r -> Sugar f ('TOr r ('TStar r))
  SStar :: Sugar f r -> Sugar f ('TStar r)
  SThen :: Sugar f l -> Sugar f r -> Sugar f ('TThen l r)

sugar :: TestEquality f => RegexRepr f r -> Sugar f r
sugar =
  \case
    Empty -> SEmpty
    Fail -> SFail
    Lit l -> SLit l
    Or l Empty -> SOptL (sugar l)
    Or Empty r -> SOptR (sugar r)
    Or (Star l) r | Just Refl <- testEquality l r -> SSomeL (sugar l)
    Or l (Star r) | Just Refl <- testEquality l r -> SSomeR (sugar l)
    Or l r -> SOr (sugar l) (sugar r)
    Star r -> SStar (sugar r)
    Then l r -> SThen (sugar l) (sugar r)

-- | Lift top-level 'Or's (i.e., those not under quantifiers)
liftOrs :: Sugar f r -> Seq (Some (Sugar f))
liftOrs =
  \case
    SEmpty -> Seq.singleton (Some SEmpty)
    SFail -> Seq.empty
    SLit x -> Seq.singleton (Some (SLit x))
    r@(SOptL {}) -> Seq.singleton (Some r)
    r@(SOptR {}) -> Seq.singleton (Some r)
    SOr l r -> liftOrs l <> liftOrs r
    r@(SSomeL {}) -> Seq.singleton (Some r)
    r@(SSomeR {}) -> Seq.singleton (Some r)
    SThen l r ->
      Seq.fromList
        [ Some (SThen x y)
        | Some x <- Foldable.toList (liftOrs l)
        , Some y <- Foldable.toList (liftOrs r)
        ]
    SStar r -> Seq.singleton (Some (SStar r))

prettySugar ::
  -- | Separator for 'SThen'
  PP.Doc ann ->
  (forall t. f t -> PP.Doc ann) ->
  Sugar f r ->
  PP.Doc ann
prettySugar sep f =
  \case
    SEmpty -> ""
    SFail -> "∅"
    SLit l -> f l
    SOptL (SLit l) -> f l PP.<> "?"
    SOptL l -> PP.parens (prettySugar sep f l) PP.<> "?"
    SOptR (SLit r) -> f r PP.<> "?"
    SOptR r -> PP.parens (prettySugar sep f r) PP.<> "?"
    SOr l r -> PP.parens (prettySugar sep f l PP.<> "|" PP.<> prettySugar sep f r)
    SSomeL (SLit l) -> f l PP.<> "+"
    SSomeL l -> PP.parens (prettySugar sep f l) PP.<> "+"
    SSomeR (SLit r) -> f r PP.<> "+"
    SSomeR r -> PP.parens (prettySugar sep f r) PP.<> "+"
    SStar (SLit l) -> f l PP.<> "*"
    SStar r -> PP.parens (prettySugar sep f r) PP.<> "*"
    SThen l r -> prettySugar sep f l PP.<> sep PP.<> prettySugar sep f r