NFA simulator based on regex type derivatives

#include "share/atspre_staload.hats"

// following: http://languagengine.co/blog/differentiating-regular-expressions/

staload _ = "prelude/DATS/list_vt.dats"

datatype Regex
  = RElit of char         // individual characters are regexes
  | REcat of (Regex, Regex)  // if E and F are regexes, so is EF
  | REalt of (Regex, Regex)  // if E and F are regexes, so is E|F
  | REopt of Regex        // if E is a regex, so is E?
  | RErep of Regex        // if E is a regex, so is E+

datatype LocalRegexContext
  = InSeqL of Regex  // a Seq with the left arg removed
  | InSeqR of Regex  // a Seq with the right arg removed
  | InAltL of Regex  // an Alt with the left arg removed
  | InAltR of Regex  // an Alt with the right arg removed
  | InOpt         // an Opt with its only arg removed
  | InRep         // a Rep with its only arg removed

typedef RegexContext = List (LocalRegexContext)

extern
fun
plug (Regex, LocalRegexContext): Regex

implement
plug (e, ctx) =
  case+ ctx of
  | InSeqL f => REcat (e, f) // e in _f is ef
  | InSeqR f => REcat (f, e) // e in f_ is fe
  | InAltL f => REalt (e, f) // e in _|f is e|f
  | InAltR f => REalt (f, e) // e in f|_ is f|e
  | InOpt () => REopt e // e in _? is e?
  | InRep () => RErep e // e in _+ is e+

extern
fun
plugContext (Regex, RegexContext): Regex
implement
plugContext (e, ctx) = let
  implement
  list_foldleft$fopr<Regex><LocalRegexContext> (acc, ctx) = plug (acc, ctx)
  val res = list_foldleft<Regex><LocalRegexContext> (ctx, e)
in
  res
end


datatype Direction = Starting | Finishing

typedef RegexState = @(Direction, Regex, RegexContext)

extern
fun
next (RegexState): List_vt RegexState

implement
next (rst) =
  case+ rst.0 of
  | Starting () => let
  in
    // this group corresponds to the edges out of starting states
    case+ (rst.1, rst.2) of
    | (RElit x, c) => $list_vt{RegexState} ((Finishing, RElit x, c) : RegexState)
    | (REcat (e, f), c) => let prval () = lemma_list_param c in $list_vt{RegexState} ((Starting, e, list_cons (InSeqL f, c))) end
    | (REalt (e, f), c) => let prval () = lemma_list_param c in $list_vt{RegexState} ((Starting, e, list_cons (InAltL f, c)), (Starting, f, list_cons (InAltR e, c))) end
    | (REopt e, c) => let prval () = lemma_list_param c in  $list_vt{RegexState} ((Starting, e, list_cons (InOpt, c)), (Finishing, REopt e, c)) end
    | (RErep e, c) => let prval () = lemma_list_param c in  $list_vt{RegexState} ((Starting, e, list_cons (InRep, c))) end
  end
  | Finishing () => (
    // this group corresponds to the edges out of finishing states
    case+ (rst.1, rst.2) of
    | (e, list_cons (InSeqL f, c)) => $list_vt{RegexState} ((Starting, f, list_cons (InSeqR e, c)))
    | (f, list_cons (InSeqR e, c)) => $list_vt{RegexState} ((Finishing, REcat (e, f), c))
    | (e, list_cons (InAltL f, c)) => $list_vt{RegexState} ((Finishing, REalt (e, f), c))
    | (f, list_cons (InAltR e, c)) => $list_vt{RegexState} ((Finishing, REalt (e, f), c))
    | (e, list_cons (InOpt (), c)) => $list_vt{RegexState} ((Finishing, REopt e, c))
    | (e, list_cons (InRep (), c)) => $list_vt{RegexState} ((Finishing, RErep e, c), (Starting, e, list_cons (InRep, c)))
    | (_, _) => list_vt_nil ()
  )

extern
fun
emit (RegexState): List_vt @(char,RegexState)
implement
emit (s) = let
  val s' = next (s)
  typedef T = @(char,RegexState)
  vtypedef VT = List_vt T
  implement
  list_vt_map$fopr<RegexState><VT> (s') =
    case+ s.0 of
    | Finishing () => (
        case+ s.1 of
        | RElit x => $list_vt{T} (@(x, s') : T)
        | _ => emit s'
      )
    | _ => emit s'
  // end of [list_vt_map$fopr]
  val res = list_vt_map<RegexState><VT> (s')
  val () = list_vt_free (s')
  val res = list_vt_concat<T> (res)
in
  res
end

extern
fun
finishesTrivially (RegexState): bool
implement
finishesTrivially (st) =
  case+ st.0 of
  | Starting () => let
      val c = st.2
      prval () = lemma_list_param (c)
    in
      case+ st.1 of
      | RElit _ => false
      | REcat (e, f) => finishesTrivially @(Starting (), e, list_cons (InSeqL f, c))
      | REalt (e, f) => finishesTrivially @(Starting (), e, list_cons (InAltL f, c))
      | REopt e => finishesTrivially @(Finishing (), REopt e, c)
      | RErep e => finishesTrivially @(Starting (), e, list_cons (InRep, c))
    end
  | Finishing () => let
      val c = st.2
      prval () = lemma_list_param (c)
    in
      case+ c of
      | list_nil () => true
      | list_cons (InSeqL f, c) => finishesTrivially @(Starting (), f, list_cons (InSeqR st.1, c))
      | list_cons (InSeqR e, c) => finishesTrivially @(Finishing (), REcat (e, st.1), c)
      | list_cons (InAltL f, c) => finishesTrivially @(Finishing (), REcat (st.1, f), c)
      | list_cons (InAltR e, c) => finishesTrivially @(Finishing (), REcat (e, st.1), c)
      | list_cons (InOpt (), c) => finishesTrivially @(Finishing (), REopt st.1, c)
      | list_cons (InRep (), c) => finishesTrivially @(Finishing (), RErep st.1, c)
    end

extern
fun
consume (char, RegexState): List_vt RegexState
implement
consume (c, s) = let
  val s' = emit (s)
  typedef T = @(char,RegexState)
  implement
  list_vt_filter$pred<T> (x) = (x.0 = c)
  val res = list_vt_filter<T> (s')
  implement
  list_vt_map$fopr<T><RegexState> (x) = x.1
  val res0 = list_vt_map<T><RegexState> (res)
  val () = list_vt_free (res)
in
  res0
end

extern
fun
recognize (string, Regex): bool
implement
recognize (cs, e) = let
  fun
  go {n:nat} (s: string(n), rs: List_vt RegexState): bool =
    if string_is_empty (s) then let
      implement
      list_vt_foreach$cont<RegexState><bool> (x, env) = ~env
      implement
      list_vt_foreach$fwork<RegexState><bool> (x, env) = env := finishesTrivially (x)
      var ans = (g0ofg1)false
      val () = list_vt_foreach_env<RegexState><bool> (rs, ans)
      val () = list_vt_free (rs)
    in
      ans
    end else let
      val c = string_head (s)
      val s = string_tail (s)
      implement
      list_vt_map$fopr<RegexState><List_vt RegexState> (s) = consume (c, s)
      val res = list_vt_map<RegexState><List_vt RegexState> (rs)
      val () = list_vt_free (rs)
      val res = list_vt_concat<RegexState> (res)
    in
      go (s, res)
    end
  // end of [go]
  val cs = (g1ofg0)cs
  val st = $list_vt{RegexState} (@(Starting (), e, list_nil ()))
in
  go (cs, st)
end // end of [recognize]

implement main0 () = let

  val () = println!("hello world!")

  // a(b|c+)d
  val re0 = REcat (REcat (RElit 'a', REalt (RElit 'b', RErep (RElit 'c'))), RElit 'd')
  val-true = recognize ("abd", re0)
  val-true = recognize ("acd", re0)
  val-true = recognize ("accd", re0)
  val-true = recognize ("acccd", re0)
  val-false = recognize ("abbd", re0)
  val-false = recognize ("efg", re0)

  val () = println!("successfully ran all tests")
in

end