Little refact and little bug fix.