5
\$\begingroup\$

This is my first Go program that I've ever written. The code can also be found here.

It allows you to create a search index and add words to it. The twist is that when searching for a word within the index you can collect all words within a given Levenshtein distance. This is done by walking multiple branches of the trie, one for each edit that is "made". For example to simulate deleting a character, we just skip/ignore the current node in the trie and keep walking, but we remember that we've consumed an edit.

Since this is my first go program I'm mostly interested in best practices regarding things like naming conventions, variable declarations, pointers and utf-8 handling using runes. I've also tried to keep an eye on performance.

Thank you very much!

package nicenshtein

import (
    "bufio"
    "os"
    "strings"
    "unicode/utf8"
)

//A trie structure that maps runes to a list of following (child-) runes.
//`word` serves two purposes:
//1. If it is not an empty string, it marks the end of a word like a flag
//2. It stores the word that the path to it spells
type RuneNode struct {
    children map[rune]*RuneNode
    word     string
}

type Nicenshtein struct {
    root *RuneNode
}

func NewNicenshtein() Nicenshtein {
    var nice Nicenshtein
    nice.root = &RuneNode{make(map[rune]*RuneNode), ""}

    return nice
}

func (nice *Nicenshtein) IndexFile(filePath string) error {
    file, err := os.Open(filePath)

    if err != nil {
        return err
    }

    defer file.Close()

    scanner := bufio.NewScanner(file)

    for scanner.Scan() {
        nextWord := strings.TrimSpace(scanner.Text())
        nice.AddWord(nextWord)
    }

    if err := scanner.Err(); err != nil {
        return err
    }

    return nil
}

func (nice *Nicenshtein) AddWord(word string) {
    if len(word) == 0 {
        return
    }

    var currentNode *RuneNode = nice.root

    for index, runeValue := range word {
        childNode, ok := currentNode.children[runeValue]

        //We have not indexed this rune yet, create a new entry.
        if !ok {
            childNode = &RuneNode{make(map[rune]*RuneNode), ""}
            currentNode.children[runeValue] = childNode
        }

        //The node at the end of a word stores the full word, which also marks the end.
        //This makes the index less memory efficient, but vastly improves query performance.
        //Otherwise each query would need to collect the runes along the path and concat the word.
        if index == len(word)-len(string(runeValue)) {
            childNode.word = word
        }

        currentNode = childNode
    }
}

func (nice *Nicenshtein) ContainsWord(word string) bool {
    var currentNode *RuneNode = nice.root

    for _, runeValue := range word {
        childNode, ok := currentNode.children[runeValue]

        //Current rune not indexed, dead end.
        if !ok {
            return false
        }

        currentNode = childNode
    }

    //Does this node terminate with the given word?
    return currentNode.word == word
}

func (nice *Nicenshtein) CollectWords(out *map[string]byte, word string, maxDistance byte) {
    nice.collectWords(out, nice.root, word, 0, maxDistance)
}

func (nice *Nicenshtein) collectWords(out *map[string]byte, currentNode *RuneNode, word string, distance byte, maxDistance byte) {
    //We have eated all runes, let's see if we have reached a node with a valid word.
    if len(word) == 0 {
        if currentNode.word != "" {
            knownDistance, ok := (*out)[currentNode.word]

            //We have not seen this word or we have found a smaller distance.
            if !ok || distance < knownDistance {
                (*out)[currentNode.word] = distance
            }
        }

        return
    }

    runeValue, size := utf8.DecodeRuneInString(word)
    wordWithoutFirstRune := word[size:]
    nextNode := currentNode.children[runeValue]

    if nextNode != nil {
        //Move forward by one rune without incrementing the distance.
        //This is just regular trie walking sans Levenshtein.
        nice.collectWords(out, nextNode, wordWithoutFirstRune, distance, maxDistance)
    }

    //Here we keep walking the trie, but with a twist.
    //We do each of the Levenshtein edits at the current position
    //and walk the trie as if nothing cool has happened.
    if distance < maxDistance {
        distance++

        //For substitution and insertion we need to apply them
        //for every rune at the current node.
        for runeValue, _ := range currentNode.children {
            //Substitution (replace the first rune with the current one).
            nice.collectWords(out, currentNode, string(runeValue)+wordWithoutFirstRune, distance, maxDistance)

            //Insertion (add the current rune as prefix).
            nice.collectWords(out, currentNode, string(runeValue)+word, distance, maxDistance)
        }

        //Deletion (skip first rune).
        nice.collectWords(out, currentNode, wordWithoutFirstRune, distance, maxDistance)
    }
}

Usage example from the tests

func TestCollectWords(t *testing.T) {
    nice := NewNicenshtein()

    nice.AddWord("Prinzhorn")
    nice.AddWord("prinzhorn")
    nice.AddWord("Crème fraîche")
    nice.AddWord("👻💩💩👻")

    closestWords := make(map[string]byte)

    nice.CollectWords(&closestWords, "Prinzhorn", 2)

    if !reflect.DeepEqual(closestWords, map[string]byte{"Prinzhorn": 0, "prinzhorn": 1}) {
        t.Error("Prinzhorn + prinzhorn not found")
    }

    closestWords = make(map[string]byte)

    nice.CollectWords(&closestWords, "Creme fraîche", 2)

    if !reflect.DeepEqual(closestWords, map[string]byte{"Crème fraîche": 1}) {
        t.Error("Crème fraîche not found")
    }

    closestWords = make(map[string]byte)

    nice.CollectWords(&closestWords, "👻💩💩💩👻", 2)

    if !reflect.DeepEqual(closestWords, map[string]byte{"👻💩💩👻": 1}) {
        t.Error("👻💩💩👻 not found")
    }
}
\$\endgroup\$
2
\$\begingroup\$

You've written a very nice package! Clean, readable and well commented code.


Just a few things to improve:

  • No need to pass Go maps by pointer. It is already a reference type just like slices are. Pass them by reference, it's 100% fine.
  • You can write map[rune]*RuneNode{} directly without make() since you don't specify the size argument.
  • for i, _ := range is same as for i := range. Don't bother to write an unnecessarily underscore, just omit it.
  • Instead of var currentNode *RuneNode = nice.root write currentNode := nice.root directly.
  • You may add a short comment with package description right before the package statement. It plays nicely with your package page at godoc.org. Also check their lint page for things to improve.
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.