Refactor ast "pattern matching" to a single function