content: package org.apache.lucene.search.suggest.jaspell;
  
  import java.io.DataInputStream;
  import java.io.DataOutputStream;
  import java.io.File;
  import java.io.FileInputStream;
  import java.io.FileOutputStream;
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.List;
  
  import org.apache.lucene.search.spell.SortedIterator;
  import org.apache.lucene.search.spell.TermFreqIterator;
  import org.apache.lucene.search.suggest.Lookup;
  import org.apache.lucene.search.suggest.UnsortedTermFreqIteratorWrapper;
  import org.apache.lucene.search.suggest.jaspell.JaspellTernarySearchTrie.TSTNode;
  
  public class JaspellLookup extends Lookup {
    JaspellTernarySearchTrie trie = new JaspellTernarySearchTrie();
    private boolean usePrefix = true;
    private int editDistance = 2;
  
    @Override
    public void build(TermFreqIterator tfit) throws IOException {
      if (tfit instanceof SortedIterator) {
        // make sure it's unsorted
        tfit = new UnsortedTermFreqIteratorWrapper(tfit);
      }
      trie = new JaspellTernarySearchTrie();
      trie.setMatchAlmostDiff(editDistance);
      while (tfit.hasNext()) {
        String key = tfit.next();
        float freq = tfit.freq();
        if (key.length() == 0) {
          continue;
        }
        trie.put(key, new Float(freq));
      }
    }
  
    @Override
    public boolean add(String key, Object value) {
      trie.put(key, value);
      // XXX
      return false;
    }
  
    @Override
    public Object get(String key) {
      return trie.get(key);
    }
  
    @Override
    public List<LookupResult> lookup(String key, boolean onlyMorePopular, int num) {
      List<LookupResult> res = new ArrayList<LookupResult>();
      List<String> list;
      int count = onlyMorePopular ? num * 2 : num;
      if (usePrefix) {
        list = trie.matchPrefix(key, count);
      } else {
        list = trie.matchAlmost(key, count);
      }
      if (list == null || list.size() == 0) {
        return res;
        
      }
      int maxCnt = Math.min(num, list.size());
      if (onlyMorePopular) {
        LookupPriorityQueue queue = new LookupPriorityQueue(num);
        for (String s : list) {
          float freq = (Float)trie.get(s);
          queue.insertWithOverflow(new LookupResult(s, freq));
        }
        for (LookupResult lr : queue.getResults()) {
          res.add(lr);
        }
      } else {
        for (int i = 0; i < maxCnt; i++) {
          String s = list.get(i);
          float freq = (Float)trie.get(s);
          res.add(new LookupResult(s, freq));
        }      
      }
      return res;
    }
  
    public static final String FILENAME = "jaspell.dat";
    private static final byte LO_KID = 0x01;
    private static final byte EQ_KID = 0x02;
    private static final byte HI_KID = 0x04;
    private static final byte HAS_VALUE = 0x08;
   
    
    @Override
    public boolean load(File storeDir) throws IOException {
      File data = new File(storeDir, FILENAME);
      if (!data.exists() || !data.canRead()) {
        return false;
      }
      DataInputStream in = new DataInputStream(new FileInputStream(data));
      TSTNode root = trie.new TSTNode('\0', null);
      try {
        readRecursively(in, root);
        trie.setRoot(root);
      } finally {
        in.close();
      }
      return true;
    }
    
    private void readRecursively(DataInputStream in, TSTNode node) throws IOException {
      node.splitchar = in.readChar();
      byte mask = in.readByte();
      if ((mask & HAS_VALUE) != 0) {
        node.data = new Float(in.readFloat());
      }
      if ((mask & LO_KID) != 0) {
        TSTNode kid = trie.new TSTNode('\0', node);
        node.relatives[TSTNode.LOKID] = kid;
        readRecursively(in, kid);
      }
      if ((mask & EQ_KID) != 0) {
        TSTNode kid = trie.new TSTNode('\0', node);
        node.relatives[TSTNode.EQKID] = kid;
        readRecursively(in, kid);
      }
      if ((mask & HI_KID) != 0) {
        TSTNode kid = trie.new TSTNode('\0', node);
        node.relatives[TSTNode.HIKID] = kid;
        readRecursively(in, kid);
      }
    }
  
    @Override
    public boolean store(File storeDir) throws IOException {
      if (!storeDir.exists() || !storeDir.isDirectory() || !storeDir.canWrite()) {
        return false;
      }
      TSTNode root = trie.getRoot();
      if (root == null) { // empty tree
        return false;
      }
      File data = new File(storeDir, FILENAME);
      DataOutputStream out = new DataOutputStream(new FileOutputStream(data));
      try {
        writeRecursively(out, root);
        out.flush();
      } finally {
        out.close();
      }
      return true;
    }
    
    private void writeRecursively(DataOutputStream out, TSTNode node) throws IOException {
      if (node == null) {
        return;
      }
      out.writeChar(node.splitchar);
      byte mask = 0;
      if (node.relatives[TSTNode.LOKID] != null) mask |= LO_KID;
      if (node.relatives[TSTNode.EQKID] != null) mask |= EQ_KID;
      if (node.relatives[TSTNode.HIKID] != null) mask |= HI_KID;
      if (node.data != null) mask |= HAS_VALUE;
      out.writeByte(mask);
      if (node.data != null) {
        out.writeFloat((Float)node.data);
      }
      writeRecursively(out, node.relatives[TSTNode.LOKID]);
      writeRecursively(out, node.relatives[TSTNode.EQKID]);
      writeRecursively(out, node.relatives[TSTNode.HIKID]);
    }
  }
  