SML modules implementing sets, NFA and DFA

Automata.sml

val _ = use "Set.sml";

functor Automata(eqtype alpha) :> 
sig
    type nfa
    type dfa
    type regex

    val Union : regex * regex -> regex
    val Concat : regex * regex -> regex
    val Star : regex -> regex
    val Null : regex
    val Emp : regex
    val Sym : alpha -> regex
    val Plus : regex -> regex
    val Question : regex -> regex
    val Times : regex * int -> regex
    val UnionList : alpha list -> regex

    val regexToNFA : regex -> nfa
    val matchesNFA : nfa -> alpha list -> bool
    val NFAtoDFA : nfa -> dfa
    val matchesDFA : dfa -> alpha list -> bool
end
= struct
    datatype lbl = Epislon | Label of alpha

    structure NFAStateSet = EqListSetFn(type t = int)
    structure NFATransSet = EqListSetFn(type t = NFAStateSet.t * lbl * NFAStateSet.t)

    structure NFAStateToTrans = CrossSet(struct structure S1 = NFAStateSet structure S2 = NFATransSet end)
    structure NFATransToState = CrossSet(struct structure S1 = NFATransSet structure S2 = NFAStateSet end)

    structure DFAStateSet = ListSetFn(struct 
        type t = NFAStateSet.t NFAStateSet.set
        fun equal(x, y) = NFAStateSet.equals x y
    end)
    structure DFATransSet = ListSetFn(struct
        type t = DFAStateSet.t * alpha * DFAStateSet.t
        fun equal((f1, l1, t1), (f2, l2, t2)) = (NFAStateSet.equals f1 f2) andalso (l1 = l2) andalso (NFAStateSet.equals t1 t2)
    end)

    structure DFAStateToTrans = CrossSet(struct structure S1 = DFAStateSet structure S2 = DFATransSet end)
    structure DFATransToState = CrossSet(struct structure S1 = DFATransSet structure S2 = DFAStateSet end)

    type nfa = {
        states : int,
        trans : NFATransSet.t NFATransSet.set,
        accept : NFAStateSet.t NFAStateSet.set
    }
    type dfa = {
        start : DFAStateSet.t,
        states : DFAStateSet.t DFAStateSet.set,
        trans : DFATransSet.t DFATransSet.set,
        accept : DFAStateSet.t DFAStateSet.set
    }
    datatype regex =
            Union of regex * regex
        |   Concat of regex * regex
        |   Star of regex
        |   Null
        |   Emp
        |   Sym of alpha

    fun Plus r = Concat(r, Star r)
    fun Question r = Union(r, Emp)
    fun Times (r, 0) = Emp
    |   Times (r, n) = Concat(r, Times(r, n - 1))
    fun UnionList [] = Emp
    |   UnionList [x] = Sym x
    |   UnionList (x::xs) = Union(Sym x, UnionList xs)

    fun regexToNFA Null = {states = 1, trans = NFATransSet.empty, accept = NFAStateSet.empty}
    |   regexToNFA Emp = {states = 1, trans = NFATransSet.empty, accept = NFAStateSet.singleton 0}
    |   regexToNFA (Sym a) = {states = 2, trans = NFATransSet.singleton (0, Label a, 1), accept = NFAStateSet.singleton 1}
    |   regexToNFA (Union (r1, r2)) =
        let
            val nfa1 = regexToNFA r1
            val nfa2 = regexToNFA r2
        in
            {
                states = 1 + #states nfa1 + #states nfa2,
                trans = NFATransSet.bigUnion
                    [
                        NFATransSet.singleton (0, Epislon, 1),
                        NFATransSet.singleton (0, Epislon, 1 + #states nfa1),
                        NFATransSet.map (fn (f, l, t) => (f + 1, l, t + 1)) (#trans nfa1),
                        NFATransSet.map (fn (f, l, t) => (f + 1 + (#states nfa1), l, t + 1 + (#states nfa1))) (#trans nfa2)
                    ],
                accept = 
                    NFAStateSet.union
                        (NFAStateSet.map (fn x => x + 1) (#accept nfa1))
                        (NFAStateSet.map (fn x => x + 1 + (#states nfa1)) (#accept nfa2))
            }
        end
    |   regexToNFA (Concat (r1, r2)) =
        let
            val nfa1 = regexToNFA r1
            val nfa2 = regexToNFA r2
        in
            {
                states = #states nfa1 + #states nfa2,
                trans = NFATransSet.bigUnion
                    [
                        #trans nfa1,
                        NFATransSet.map (fn (f, l, t) => (f + (#states nfa1), l, t + (#states nfa1))) (#trans nfa2),
                        NFAStateToTrans.map (fn x => (x, Epislon, #states nfa1)) (#accept nfa1)
                    ],
                accept = NFAStateSet.map (fn x => x + (#states nfa1)) (#accept nfa2)
            }
        end
    |   regexToNFA (Star r) =
        let
            val nfa = regexToNFA r
        in
            {
                states = 1 + #states nfa,
                trans = NFATransSet.bigUnion
                    [
                        NFATransSet.map (fn (f, l, t) => (f + 1, l, t + 1)) (#trans nfa),
                        NFAStateToTrans.map (fn x => (x + 1, Epislon, 0)) (#accept nfa),
                        NFATransSet.singleton (0, Epislon, 1)
                    ],
                accept = NFAStateSet.singleton 0
            }
        end

    fun epsilonClosureFromState nfa state =
    let
        val visited = ref (NFAStateSet.singleton state)
        fun helper s =
        let
            val directEpsilon = NFATransToState.map (fn (_, _, t) => t) (NFATransSet.select (fn (f, x, t) => f = s andalso x = Epislon andalso (not (NFAStateSet.member t (!visited)))) (#trans nfa))
            val _ = visited := NFAStateSet.union (!visited) directEpsilon
            val _ = NFAStateSet.app helper directEpsilon
        in
            ()
        end
        val _ = helper state
    in
        !visited
    end

    fun epsilonClosureFromStates nfa states = NFAStateSet.bigUnion (map (epsilonClosureFromState nfa) (NFAStateSet.toList states))

    fun move nfa symbol states = NFATransToState.map (fn (_, _, t) => t) (NFATransSet.select (fn (f, x, _) => NFAStateSet.member f states andalso x = Label symbol) (#trans nfa))

    fun reachableFromState nfa s symbol = epsilonClosureFromStates nfa (move nfa symbol (epsilonClosureFromState nfa s))

    fun matchesNFA nfa lst =
    let
        fun curStates [] s = s
        |   curStates (x::xs) s =
                curStates xs (NFAStateSet.bigUnion (map (fn st => reachableFromState nfa st x) (NFAStateSet.toList s)))
    in
        NFAStateSet.size (NFAStateSet.intersect (curStates lst (NFAStateSet.singleton 0)) (#accept nfa)) > 0
    end

    structure LabelSet = EqListSetFn(type t = lbl)
    structure AlphaSet = EqListSetFn(type t = alpha)
    structure TransToLabel = CrossSet(struct structure S1 = NFATransSet structure S2 = LabelSet end)
    structure LabelToAlpha = CrossSet(struct structure S1 = LabelSet structure S2 = AlphaSet end)

    fun getNFAalphabet nfa = LabelToAlpha.map (fn Label x => x | _ => raise Fail "placeholder") (LabelSet.select (fn Label _ => true | _ => false) (TransToLabel.map (fn (_, l, _) => l) (#trans nfa)))

    structure AlphaStatesSet = ListSetFn(struct
        type t = alpha * NFAStateSet.t NFAStateSet.set
        fun equal ((a1, s1), (a2, s2)) = (a1 = a2) andalso NFAStateSet.equals s1 s2
    end)
    structure AlphaToAlphaStates = CrossSet(struct structure S1 = AlphaSet structure S2 = AlphaStatesSet end)
    structure AlphaStatesToDFAState = CrossSet(struct structure S1 = AlphaStatesSet structure S2 = DFAStateSet end)
    structure AlphaStatesToDFATrans = CrossSet(struct structure S1 = AlphaStatesSet structure S2 = DFATransSet end)

    fun NFAtoDFA nfa =
    let
        val alphabet = getNFAalphabet nfa

        val start = epsilonClosureFromState nfa 0

        fun helper [] states trans = (states, trans)
        |   helper (cur::left) states trans =
            if DFAStateSet.member cur states then
                helper left states trans
            else
            let
                val t = AlphaToAlphaStates.map (fn a => (a, epsilonClosureFromStates nfa (move nfa a cur))) alphabet
                val unexplored = AlphaStatesToDFAState.map (fn (_, x) => x) t
                val new_trans = AlphaStatesToDFATrans.map (fn (a, s) => (cur, a, s)) t
            in
                helper (left @ (DFAStateSet.toList unexplored)) (DFAStateSet.add cur states) (DFATransSet.union trans new_trans)
            end

        val (states, trans) = helper [start] DFAStateSet.empty DFATransSet.empty
        val accept = DFAStateSet.select (NFAStateSet.exists (fn s => NFAStateSet.member s (#accept nfa))) states
    in
        {
            start = start,
            states = states,
            trans = trans,
            accept = accept
        }
    end

    fun matchesDFA dfa lst =
    let
        fun curState [] s = s
        |   curState (x::xs) s =
                curState xs (DFAStateSet.pick (DFATransToState.map (fn (_, _, t) => t) (DFATransSet.select (fn (f, l, _) => NFAStateSet.equals f s andalso l = x) (#trans dfa))))
    in
        DFAStateSet.member (curState lst (#start dfa)) (#accept dfa)
            handle Subscript => false
    end
end

Set.sml

signature SET = sig
    type t
    type 'a set

    val empty : 'a set
    val singleton : 'a -> 'a set

    val size : 'a set -> int

    val member : t -> t set -> bool
    val subset : t set -> t set -> bool
    val equals : t set -> t set -> bool

    val add : t -> t set -> t set
    val drop : t -> t set -> t set
    val pick : 'a set -> 'a
    val pickdrop : 'a set -> 'a * 'a set

    val union : t set -> t set -> t set
    val intersect : t set -> t set -> t set
    val difference : t set -> t set -> t set

    val bigUnion : t set list -> t set
    val bigIntersect : t set list -> t set

    val fromList : t list -> t set
    val toList : 'a set -> 'a list

    val select : ('a -> bool) -> 'a set -> 'a set
    val map : ('a -> t) -> 'a set -> t set
    val app : ('a -> unit) -> 'a set -> unit
    val exists : ('a -> bool) -> 'a set -> bool
end

functor CrossSet (structure S1 : SET structure S2 : SET) :>
sig
    val convert : S2.t S1.set -> S2.t S2.set
    val map : ('a -> S2.t) -> 'a S1.set -> S2.t S2.set
end
= struct
    fun convert s = S2.fromList (S1.toList s)

    fun map f = S2.fromList o (List.map f) o S1.toList
end

functor ListSetFn (EQ : sig type t val equal : t * t -> bool end) :> SET where type t = EQ.t = struct
    type t = EQ.t
    type 'a set = 'a list

    val empty = []
    fun singleton x = [x]

    val size = List.length

    fun member x l = List.exists (fn y => EQ.equal(x, y)) l

    fun subset [] _ = true
    |   subset (e::es) s = member e s andalso subset es s

    fun equals s1 s2 = size s1 = size s2 andalso subset s1 s2

    fun add x [] = [x]
    |   add x (e::es) = if EQ.equal(e, x) then (e::es) else e :: (add x es)

    fun drop x [] = []
    |   drop x (e::es) = if EQ.equal(e, x) then es else e :: (drop x es)

    fun pick [] = raise Subscript
    |   pick (x::_) = x

    fun pickdrop [] = raise Subscript
    |   pickdrop (x::xs) = (x, xs)

    fun union [] s = s
    |   union (x::xs) s = union xs (add x s)

    fun intersect [] _ = []
    |   intersect (e::es) s = if member e s then e :: (intersect es s) else intersect es s

    fun difference [] _ = []
    |   difference (e::es) s = if member e s then difference es s else e :: (difference es s)

    fun bigUnion [] = []
    |   bigUnion (s::ss) = union s (bigUnion ss)

    fun bigIntersect [] = []
    |   bigIntersect (s::ss) = intersect s (bigIntersect ss)

    fun fromList l = bigUnion (map (fn x => [x]) l)
    fun toList l = l

    val select = List.filter
    fun map f = fromList o (List.map f)
    val app = List.app
    val exists = List.exists
end

functor EqListSetFn(eqtype t) :> SET where type t = t = struct
    structure T = ListSetFn(struct type t = t val equal = op= end)
    open T
end