int _yr_ac_optimize_failure_links(
YR_AC_AUTOMATON* automaton)
{
QUEUE queue = { NULL, NULL};
// Push root's children.
YR_AC_STATE* root_state = automaton->root;
YR_AC_STATE* state = root_state->first_child;
while (state != NULL)
{
FAIL_ON_ERROR(_yr_ac_queue_push(&queue, state));
state = state->siblings;
}
while (!_yr_ac_queue_is_empty(&queue))
{
YR_AC_STATE* current_state = _yr_ac_queue_pop(&queue);
if (current_state->failure != root_state)
{
if (_yr_ac_transitions_subset(current_state, current_state->failure))
current_state->failure = current_state->failure->failure;
}
// Push childrens of current_state
state = current_state->first_child;
while (state != NULL)
{
FAIL_ON_ERROR(_yr_ac_queue_push(&queue, state));
state = state->siblings;
}
}
return ERROR_SUCCESS;
}
void yr_ac_create_failure_links(
YR_ARENA* arena,
YR_AC_AUTOMATON* automaton)
{
YR_AC_STATE_TRANSITION transition;
YR_AC_STATE* current_state;
YR_AC_STATE* failure_state;
YR_AC_STATE* temp_state;
YR_AC_STATE* state;
YR_AC_STATE* transition_state;
YR_AC_STATE* root_state;
YR_AC_MATCH* match;
QUEUE queue;
queue.head = NULL;
queue.tail = NULL;
root_state = automaton->root;
// Set the failure link of root state to itself.
root_state->failure = root_state;
// Push root's children and set their failure link to root.
state = _yr_ac_first_transition(root_state, &transition);
while (state != NULL)
{
_yr_ac_queue_push(&queue, state);
state->failure = root_state;
state = _yr_ac_next_transition(root_state, &transition);
}
// Traverse the trie in BFS order calculating the failure link
// for each state.
while (!_yr_ac_queue_is_empty(&queue))
{
current_state = _yr_ac_queue_pop(&queue);
match = current_state->matches;
if (match != NULL)
{
while (match->next != NULL)
match = match->next;
if (match->backtrack > 0)
match->next = root_state->matches;
}
else
{
current_state->matches = root_state->matches;
}
transition_state = _yr_ac_first_transition(
current_state,
&transition);
while (transition_state != NULL)
{
_yr_ac_queue_push(&queue, transition_state);
failure_state = current_state->failure;
while (1)
{
temp_state = yr_ac_next_state(
failure_state,
transition.input);
if (temp_state != NULL)
{
transition_state->failure = temp_state;
if (transition_state->matches == NULL)
{
transition_state->matches = temp_state->matches;
}
else
{
match = transition_state->matches;
while (match != NULL && match->next != NULL)
match = match->next;
match->next = temp_state->matches;
}
break;
}
else
{
if (failure_state == root_state)
{
transition_state->failure = root_state;
break;
}
else
{
failure_state = failure_state->failure;
}
}
} // while(1)
transition_state = _yr_ac_next_transition(
current_state,
&transition);
}
} // while(!__yr_ac_queue_is_empty(&queue))
}
int _yr_ac_build_transition_table(
YR_AC_AUTOMATON* automaton)
{
YR_AC_STATE* state;
YR_AC_STATE* child_state;
YR_AC_STATE* root_state = automaton->root;
uint32_t slot;
QUEUE queue = { NULL, NULL};
automaton->tables_size = 1024;
automaton->t_table = (YR_AC_TRANSITION_TABLE) yr_malloc(
automaton->tables_size * sizeof(YR_AC_TRANSITION));
automaton->m_table = (YR_AC_MATCH_TABLE) yr_malloc(
automaton->tables_size * sizeof(YR_AC_MATCH_TABLE_ENTRY));
if (automaton->t_table == NULL || automaton->m_table == NULL)
{
yr_free(automaton->t_table);
yr_free(automaton->m_table);
return ERROR_INSUFFICIENT_MEMORY;
}
memset(automaton->t_table, 0,
automaton->tables_size * sizeof(YR_AC_TRANSITION));
memset(automaton->m_table, 0,
automaton->tables_size * sizeof(YR_AC_MATCH_TABLE_ENTRY));
automaton->t_table[0] = YR_AC_MAKE_TRANSITION(0, 0, YR_AC_USED_FLAG);
automaton->m_table[0].match = root_state->matches;
// Index 0 is for root node. Unused indexes start at 1.
automaton->t_table_unused_candidate = 1;
child_state = root_state->first_child;
while (child_state != NULL)
{
child_state->t_table_slot = child_state->input + 1;
automaton->t_table[child_state->input + 1] = YR_AC_MAKE_TRANSITION(
0, child_state->input + 1, YR_AC_USED_FLAG);
FAIL_ON_ERROR(_yr_ac_queue_push(&queue, child_state));
child_state = child_state->siblings;
}
while (!_yr_ac_queue_is_empty(&queue))
{
state = _yr_ac_queue_pop(&queue);
FAIL_ON_ERROR(_yr_ac_find_suitable_transition_table_slot(
automaton,
state,
&slot));
automaton->t_table[state->t_table_slot] |= ((uint64_t) slot << 32);
state->t_table_slot = slot;
automaton->t_table[slot] = YR_AC_MAKE_TRANSITION(
state->failure->t_table_slot, 0, YR_AC_USED_FLAG);
automaton->m_table[slot].match = state->matches;
// Push childrens of current_state
child_state = state->first_child;
while (child_state != NULL)
{
child_state->t_table_slot = slot + child_state->input + 1;
automaton->t_table[child_state->t_table_slot] = YR_AC_MAKE_TRANSITION(
0, child_state->input + 1, YR_AC_USED_FLAG);
FAIL_ON_ERROR(_yr_ac_queue_push(&queue, child_state));
child_state = child_state->siblings;
}
}
return ERROR_SUCCESS;
}
