Add trace function and fix some bugs